rtx-csm: WavLM transformer encoder with gated rel-pos attention
Phase 5b — full implementation of WavLMEncoderLayer.forward, replacing the Phase 5a stub. With random weights the encoder is no longer a no-op (verified: L2(output - input) > 1e-6 in unit test). - relative_position_bucket helper: T5-style bidirectional bucketing matching HF _relative_positions_bucket. 320 buckets, 800 max distance, half/half split with linear inner / log-spaced outer. - WavLmEncoderLayer::compute_position_bias (layer 0 only): builds (T,T) bucket index, embedding-looks up rel_attn_embed, permutes to (num_heads, T, T) matching HF's compute_bias output. - WavLmEncoderLayer::gated_position_bias: HF gating math verbatim — Linear(head_dim → 8), reshape (..., 2, 4) sum, sigmoid, chunk to gate_a/gate_b, compute gate_a * (gate_b * gru_rel_pos_const - 1) + 2, broadcast-multiply position_bias. - WavLmEncoderLayer::attention: multi-head self-attention with the gated bias added to scores before softmax. Standard 1/sqrt(d) scale. - WavLmEncoderLayer::forward_with_bias returns (output, position_bias) so Encoder::forward_all_layers can thread bias from layer 0 through layers 1-11 (HF's has_relative_position_bias=(i==0) pattern). - 3 new tests bring wavlm_sv to 12 tests; 75 lib tests total green. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -63,7 +63,7 @@
|
||||
use crate::error::{CsmError, Result};
|
||||
use candle_core::{DType, Device, IndexOp, Module, Tensor, D};
|
||||
use candle_nn::{
|
||||
conv1d, group_norm, layer_norm, linear, ops, Activation, Conv1d, Conv1dConfig, GroupNorm,
|
||||
conv1d, group_norm, layer_norm, linear, ops, Conv1d, Conv1dConfig, GroupNorm,
|
||||
LayerNorm, LayerNormConfig, Linear, VarBuilder,
|
||||
};
|
||||
use std::path::Path;
|
||||
@@ -243,12 +243,30 @@ impl Module for PosConv {
|
||||
}
|
||||
}
|
||||
|
||||
// -- 4. Transformer encoder layer (STUB — Phase 5b) -----------------------
|
||||
// -- 4. Transformer encoder layer with gated relative-position attention --
|
||||
|
||||
/// T5-style relative-position bucket. Maps a relative offset (k - q) into
|
||||
/// `[0, num_buckets)`. Half the buckets cover negative offsets, half cover
|
||||
/// positive; within each half, the first `max_exact = num_buckets/4` are
|
||||
/// linear, the rest log-spaced up to `max_distance`.
|
||||
fn relative_position_bucket(rel_pos: i64, num_buckets: usize, max_distance: usize) -> u32 {
|
||||
let half = num_buckets / 2;
|
||||
let mut bucket = if rel_pos > 0 { half } else { 0 };
|
||||
let abs_pos = rel_pos.unsigned_abs() as usize;
|
||||
let max_exact = half / 2;
|
||||
if abs_pos < max_exact {
|
||||
bucket += abs_pos;
|
||||
} else {
|
||||
let log_ratio = (abs_pos as f64 / max_exact as f64).ln();
|
||||
let log_factor = (max_distance as f64 / max_exact as f64).ln();
|
||||
let log_bucket = (log_ratio / log_factor) * (half - max_exact) as f64;
|
||||
let large = max_exact + log_bucket as usize;
|
||||
bucket += large.min(half - 1);
|
||||
}
|
||||
bucket as u32
|
||||
}
|
||||
|
||||
/// One WavLM encoder layer. Post-norm: residual+attn → LN → FFN-residual → final-LN.
|
||||
///
|
||||
/// **Stub**: forward returns input unchanged. Phase 5b will implement
|
||||
/// gated relative-position bias attention + 768→3072→768 FFN.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WavLmEncoderLayer {
|
||||
// Phase 5b — these will be populated:
|
||||
@@ -311,10 +329,128 @@ impl WavLmEncoderLayer {
|
||||
})
|
||||
}
|
||||
|
||||
/// **STUB**: returns input unchanged. Phase 5b implements full
|
||||
/// attention + FFN. Shape `(B, T, 768)` in/out.
|
||||
/// Compute the bucketed relative-position bias `(num_heads, T, T)` for
|
||||
/// a given query length. Only callable on layer 0 (the layer that owns
|
||||
/// `rel_attn_embed`). Result is reused by the rest of the stack.
|
||||
pub fn compute_position_bias(&self, t: usize) -> candle_core::Result<Tensor> {
|
||||
let embed = self
|
||||
.rel_attn_embed
|
||||
.as_ref()
|
||||
.ok_or_else(|| candle_core::Error::Msg(
|
||||
"compute_position_bias called on layer without rel_attn_embed (only layer 0 has one)".into(),
|
||||
))?;
|
||||
// Build (T, T) bucket index tensor on the host.
|
||||
let mut buckets = Vec::with_capacity(t * t);
|
||||
for q in 0..t {
|
||||
for k in 0..t {
|
||||
let rel = k as i64 - q as i64;
|
||||
buckets.push(relative_position_bucket(rel, REL_NUM_BUCKETS, REL_MAX_DISTANCE));
|
||||
}
|
||||
}
|
||||
let device = embed.embeddings().device();
|
||||
let idx = Tensor::from_vec(buckets, (t, t), device)?;
|
||||
// Embedding lookup: (T, T) → (T, T, num_heads). Permute to
|
||||
// (num_heads, T, T) to match the HF compute_bias output.
|
||||
let values = embed.forward(&idx)?;
|
||||
values.permute((2, 0, 1))?.contiguous()
|
||||
}
|
||||
|
||||
/// Compute the gated bias to add to attention scores.
|
||||
/// `position_bias`: `(num_heads, T, T)`; `xs`: `(B, T, embed_dim)`.
|
||||
/// Output: `(B, num_heads, T, T)` — the per-batch, per-head gated bias.
|
||||
fn gated_position_bias(
|
||||
&self,
|
||||
position_bias: &Tensor,
|
||||
xs: &Tensor,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let (b, t, _) = xs.dims3()?;
|
||||
// (B, T, embed_dim) → (B, T, num_heads, head_dim) → (B, num_heads, T, head_dim)
|
||||
let h = xs
|
||||
.reshape((b, t, NUM_HEADS, HEAD_DIM))?
|
||||
.permute((0, 2, 1, 3))?
|
||||
.contiguous()?;
|
||||
// Linear(head_dim → 8) → (B, num_heads, T, 8)
|
||||
let proj = h.apply(&self.gru_rel_pos_linear)?;
|
||||
// (B, num_heads, T, 2, 4) → sum(-1) → (B, num_heads, T, 2)
|
||||
let proj = proj.reshape((b, NUM_HEADS, t, 2, 4))?.sum(D::Minus1)?;
|
||||
let gates = ops::sigmoid(&proj)?; // (B, num_heads, T, 2)
|
||||
let gate_a = gates.narrow(D::Minus1, 0, 1)?; // (B, num_heads, T, 1)
|
||||
let gate_b = gates.narrow(D::Minus1, 1, 1)?; // (B, num_heads, T, 1)
|
||||
// gate_output = gate_a * (gate_b * const - 1) + 2
|
||||
let const_g = self
|
||||
.gru_rel_pos_const
|
||||
.broadcast_as((b, NUM_HEADS, t, 1))?
|
||||
.to_dtype(gate_b.dtype())?;
|
||||
let inner = ((gate_b * const_g)? - 1.0f64)?;
|
||||
let gate_out = ((gate_a * inner)? + 2.0f64)?; // (B, num_heads, T, 1)
|
||||
// Broadcast position_bias (num_heads, T, T) → (1, num_heads, T, T) → (B, num_heads, T, T)
|
||||
let bias = position_bias
|
||||
.unsqueeze(0)?
|
||||
.broadcast_as((b, NUM_HEADS, t, t))?
|
||||
.to_dtype(gate_out.dtype())?;
|
||||
gate_out.broadcast_mul(&bias)
|
||||
}
|
||||
|
||||
/// Multi-head self-attention with the gated bias added to the score
|
||||
/// matrix before softmax.
|
||||
fn attention(&self, xs: &Tensor, gated_bias: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let (b, t, _) = xs.dims3()?;
|
||||
let split_heads = |proj: Tensor| -> candle_core::Result<Tensor> {
|
||||
proj.reshape((b, t, NUM_HEADS, HEAD_DIM))?
|
||||
.permute((0, 2, 1, 3))?
|
||||
.contiguous()
|
||||
};
|
||||
let q = split_heads(xs.apply(&self.q_proj)?)?;
|
||||
let k = split_heads(xs.apply(&self.k_proj)?)?;
|
||||
let v = split_heads(xs.apply(&self.v_proj)?)?;
|
||||
let scale = (HEAD_DIM as f64).powf(-0.5);
|
||||
let scores = (q.matmul(&k.transpose(2, 3)?.contiguous()?)? * scale)?;
|
||||
let scores = (scores + gated_bias)?;
|
||||
let attn = ops::softmax_last_dim(&scores)?;
|
||||
let out = attn.matmul(&v)?; // (B, num_heads, T, head_dim)
|
||||
let out = out
|
||||
.permute((0, 2, 1, 3))?
|
||||
.contiguous()?
|
||||
.reshape((b, t, HIDDEN_DIM))?;
|
||||
out.apply(&self.out_proj)
|
||||
}
|
||||
|
||||
/// Forward pass. Returns `(hidden_states, position_bias)` so the
|
||||
/// encoder can thread the bias through subsequent layers.
|
||||
///
|
||||
/// `in_bias` is `Some` for layers 1..N — the bias computed by layer 0.
|
||||
/// `None` is acceptable on layer 0 (we'll compute it locally) and an
|
||||
/// error on any other layer (caller's bug).
|
||||
pub fn forward_with_bias(
|
||||
&self,
|
||||
xs: &Tensor,
|
||||
in_bias: Option<&Tensor>,
|
||||
) -> candle_core::Result<(Tensor, Tensor)> {
|
||||
let owned;
|
||||
let bias = match in_bias {
|
||||
Some(b) => b,
|
||||
None => {
|
||||
let t = xs.dim(D::Minus2)?;
|
||||
owned = self.compute_position_bias(t)?;
|
||||
&owned
|
||||
}
|
||||
};
|
||||
let gated = self.gated_position_bias(bias, xs)?;
|
||||
let attn_out = self.attention(xs, &gated)?;
|
||||
let h = (xs + attn_out)?;
|
||||
let h = h.apply(&self.attn_norm)?;
|
||||
// FFN: 768 → 3072 → GELU → 768
|
||||
let ffn = h.apply(&self.fc1)?.gelu()?.apply(&self.fc2)?;
|
||||
let h = (h + ffn)?;
|
||||
let h = h.apply(&self.final_norm)?;
|
||||
Ok((h, bias.clone()))
|
||||
}
|
||||
|
||||
/// Convenience used by the smoke test where the caller doesn't carry a
|
||||
/// bias — equivalent to layer 0's `forward_with_bias(xs, None)`.
|
||||
pub fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
Ok(xs.clone())
|
||||
let (out, _bias) = self.forward_with_bias(xs, None)?;
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -345,14 +481,20 @@ impl Encoder {
|
||||
|
||||
/// Input `(B, T, 768)`; output is `(B, T, 768)` per layer + initial,
|
||||
/// returned as a `Vec<Tensor>` of length `NUM_LAYERS + 1 = 13`.
|
||||
/// The first hidden state is the post-norm input embedding (matches
|
||||
/// the HF `output_hidden_states` convention used by `WavLMForXVector`).
|
||||
pub fn forward_all_layers(&self, xs: &Tensor) -> candle_core::Result<Vec<Tensor>> {
|
||||
let pos = self.pos_conv.forward(xs)?;
|
||||
let mut h = (xs + &pos)?;
|
||||
h = h.apply(&self.layer_norm)?;
|
||||
let mut hidden_states = Vec::with_capacity(NUM_LAYERS + 1);
|
||||
hidden_states.push(h.clone());
|
||||
// Layer 0 computes the position_bias (shared across all 12 layers).
|
||||
let mut position_bias: Option<Tensor> = None;
|
||||
for layer in &self.layers {
|
||||
h = layer.forward(&h)?;
|
||||
let (next_h, bias) = layer.forward_with_bias(&h, position_bias.as_ref())?;
|
||||
h = next_h;
|
||||
position_bias = Some(bias);
|
||||
hidden_states.push(h.clone());
|
||||
}
|
||||
Ok(hidden_states)
|
||||
@@ -692,6 +834,50 @@ mod tests {
|
||||
assert_eq!(emb.dims(), &[1, EMBEDDING_DIM]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn relative_position_bucket_basics() {
|
||||
// Distance 0 → bucket 0 (small, left side).
|
||||
assert_eq!(relative_position_bucket(0, 320, 800), 0);
|
||||
// Tiny positive → goes to right half (bucket >= 160).
|
||||
let b1 = relative_position_bucket(1, 320, 800);
|
||||
assert!(b1 >= 160 && b1 < 320);
|
||||
// Large positive → still on right half, capped to half - 1.
|
||||
let b_large = relative_position_bucket(10_000, 320, 800);
|
||||
assert_eq!(b_large, 320 - 1);
|
||||
// Negative on the left half.
|
||||
let b_neg = relative_position_bucket(-1, 320, 800);
|
||||
assert!(b_neg < 160);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encoder_layer_runs_with_layer0_bias() {
|
||||
let device = Device::Cpu;
|
||||
let (_vm, vb) = random_vb(&device);
|
||||
let layer = WavLmEncoderLayer::new(0, vb).unwrap();
|
||||
let xs = Tensor::randn(0f32, 1f32, (1, 24, HIDDEN_DIM), &device).unwrap();
|
||||
let (out, bias) = layer.forward_with_bias(&xs, None).unwrap();
|
||||
assert_eq!(out.dims(), &[1, 24, HIDDEN_DIM]);
|
||||
assert_eq!(bias.dims(), &[NUM_HEADS, 24, 24]);
|
||||
// Verify the layer actually transformed the input (not a no-op).
|
||||
let diff = (&out - &xs).unwrap();
|
||||
let l2 = diff.sqr().unwrap().sum_all().unwrap();
|
||||
let l2: f32 = l2.to_dtype(DType::F32).unwrap().to_scalar().unwrap();
|
||||
assert!(l2 > 1e-6, "encoder layer is a no-op (l2 = {l2})");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encoder_threads_position_bias_through_stack() {
|
||||
let device = Device::Cpu;
|
||||
let (_vm, vb) = random_vb(&device);
|
||||
let enc = Encoder::new(vb).unwrap();
|
||||
let xs = Tensor::randn(0f32, 1f32, (1, 24, HIDDEN_DIM), &device).unwrap();
|
||||
let states = enc.forward_all_layers(&xs).unwrap();
|
||||
assert_eq!(states.len(), NUM_LAYERS + 1);
|
||||
for s in &states {
|
||||
assert_eq!(s.dims(), &[1, 24, HIDDEN_DIM]);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cosine_similarity_self_is_one() {
|
||||
let device = Device::Cpu;
|
||||
|
||||
Reference in New Issue
Block a user