steeldb/

hrm.rs

//! **HRM two-timescale reasoning core** — a faithful candle port of the reference
//! `python/splade/hrm_core.py` + `spo_tagger.py::HRMTagger`.
//!
//! This is the architecture the reference actually uses, and it is not a stack of transformer layers:
//!
//! * the backbone is BERT's **embeddings only** (word + position + type) — no encoder. With a 128-dim,
//!   30522-token vocabulary that is ~3.9M parameters, which is where the "4M model" comes from;
//! * on top sits [`HrmCore`], two coupled recurrent timescales — a **fast** `L` stream that iterates
//!   `T_inner` times against the injected input, and a **slow** `H` stream that *steps back* once per
//!   cycle to integrate what `L` produced. Depth comes from `N_cycles × T_inner` refinement passes
//!   rather than from more weights;
//! * gradients flow only through the **final** L+H step (the reference wraps the earlier cycles in
//!   `no_grad` and detaches the carry). That one-step gradient approximation is what makes recurrent
//!   depth affordable, and it is reproduced here with `detach()`.
//!
//! Heads are **independent** linear maps off the shared refined state: BIO span typing, the epistemic
//! reading, and (via [`HrmTagger::hidden`]) span pooling for the biaffine relation head. Adding a head
//! costs one matrix, not another encoder.
//!
//! Faithfulness notes: `Block` is pre-norm `RMSNorm → MHA → RMSNorm → SwiGLU`, matching the reference.
//! Dropout is omitted — the reference uses 0.1 on the residual branches, which regularises a long
//! training run but changes no shapes or values at eval; it is the one deliberate deviation.

use candle_core::{DType, Device, Module, Result, Tensor, D};
use candle_nn::{ops::softmax, Linear, VarBuilder};

/// `x * rsqrt(mean(x²) + eps) * w` — the reference's RMSNorm (no mean subtraction, learned gain).
pub struct RmsNorm {
    w: Tensor,
    eps: f64,
}

impl RmsNorm {
    pub fn new(h: usize, vb: VarBuilder) -> Result<Self> {
        Ok(RmsNorm { w: vb.get(h, "w")?, eps: 1e-6 })
    }
    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let ms = x.sqr()?.mean_keepdim(D::Minus1)?;
        let scaled = x.broadcast_div(&(ms + self.eps)?.sqrt()?)?;
        scaled.broadcast_mul(&self.w)
    }
}

/// Multi-head self-attention with an additive key-padding mask.
struct SelfAttention {
    q: Linear,
    k: Linear,
    v: Linear,
    o: Linear,
    heads: usize,
    head_dim: usize,
}

impl SelfAttention {
    fn new(h: usize, heads: usize, vb: VarBuilder) -> Result<Self> {
        Ok(SelfAttention {
            q: candle_nn::linear(h, h, vb.pp("q"))?,
            k: candle_nn::linear(h, h, vb.pp("k"))?,
            v: candle_nn::linear(h, h, vb.pp("v"))?,
            o: candle_nn::linear(h, h, vb.pp("o"))?,
            heads,
            head_dim: h / heads,
        })
    }

    /// `x [B,T,h]`, `mask [B,T]` (1 = real token) → `[B,T,h]`.
    fn forward(&self, x: &Tensor, mask: Option<&Tensor>) -> Result<Tensor> {
        let (b, t, h) = x.dims3()?;
        let split = |p: &Linear, x: &Tensor| -> Result<Tensor> {
            p.forward(x)?.reshape((b, t, self.heads, self.head_dim))?.transpose(1, 2)?.contiguous()
        };
        let q = split(&self.q, x)?;
        let k = split(&self.k, x)?;
        let v = split(&self.v, x)?;
        let scale = 1.0 / (self.head_dim as f64).sqrt();
        let mut att = (q.matmul(&k.transpose(2, 3)?.contiguous()?)? * scale)?; // [B,heads,T,T]
        if let Some(m) = mask {
            // Additive mask: 0 for real tokens, a large negative for padding, broadcast over heads and
            // query rows. A *finite* penalty is deliberate — scaling the inverted mask by -inf computes
            // `0 * -inf` at every real position, which is NaN and poisons the whole attention matrix.
            const MASK_PENALTY: f64 = -1e9;
            let neg = m
                .to_dtype(DType::F32)?
                .affine(-1.0, 1.0)? // real 1 → 0, pad 0 → 1
                .affine(MASK_PENALTY, 0.0)? // real → 0, pad → -1e9
                .reshape((b, 1, 1, t))?;
            att = att.broadcast_add(&neg)?;
        }
        let att = softmax(&att, D::Minus1)?;
        let out = att.matmul(&v)?.transpose(1, 2)?.reshape((b, t, h))?;
        self.o.forward(&out)
    }
}

/// Pre-norm block: `RMSNorm → MHA → residual → RMSNorm → SwiGLU → residual`.
struct Block {
    n1: RmsNorm,
    attn: SelfAttention,
    n2: RmsNorm,
    w_in: Linear,
    w_out: Linear,
}

impl Block {
    fn new(h: usize, heads: usize, mult: usize, vb: VarBuilder) -> Result<Self> {
        Ok(Block {
            n1: RmsNorm::new(h, vb.pp("n1"))?,
            attn: SelfAttention::new(h, heads, vb.pp("attn"))?,
            n2: RmsNorm::new(h, vb.pp("n2"))?,
            // w_in emits 2*mult*h so it can be chunked into the SwiGLU gate pair
            w_in: candle_nn::linear(h, 2 * mult * h, vb.pp("w_in"))?,
            w_out: candle_nn::linear(mult * h, h, vb.pp("w_out"))?,
        })
    }

    fn forward(&self, x: &Tensor, mask: Option<&Tensor>) -> Result<Tensor> {
        let a = self.attn.forward(&self.n1.forward(x)?, mask)?;
        let x = (x + a)?;
        let g = self.n2.forward(&x)?;
        let uv = self.w_in.forward(&g)?;
        let half = uv.dim(D::Minus1)? / 2;
        let u = uv.narrow(D::Minus1, 0, half)?;
        let v = uv.narrow(D::Minus1, half, half)?;
        let gated = (u.silu()? * v)?;
        x + self.w_out.forward(&gated)?
    }
}

/// The two-timescale core. `L` iterates fast against the injected input; `H` steps back once per cycle to
/// integrate `L`'s state. Output is `zH + reps`, a residual refinement of the embeddings.
pub struct HrmCore {
    inj: Linear,
    l_blocks: Vec<Block>,
    h_blocks: Vec<Block>,
    z_l0: Tensor,
    z_h0: Tensor,
    pub t_inner: usize,
    pub n_cycles: usize,
}

impl HrmCore {
    pub fn new(h: usize, layers_l: usize, layers_h: usize, t_inner: usize, n_cycles: usize, heads: usize, vb: VarBuilder) -> Result<Self> {
        let l_blocks = (0..layers_l).map(|i| Block::new(h, heads, 2, vb.pp(format!("L{i}")))).collect::<Result<Vec<_>>>()?;
        let h_blocks = (0..layers_h).map(|i| Block::new(h, heads, 2, vb.pp(format!("H{i}")))).collect::<Result<Vec<_>>>()?;
        Ok(HrmCore {
            inj: candle_nn::linear(h, h, vb.pp("inj"))?,
            l_blocks,
            h_blocks,
            z_l0: vb.get(h, "z_l0")?,
            z_h0: vb.get(h, "z_h0")?,
            t_inner,
            n_cycles,
        })
    }

    fn l_step(&self, z_l: &Tensor, z_h: &Tensor, x: &Tensor, mask: Option<&Tensor>) -> Result<Tensor> {
        let mut y = ((z_l + z_h)? + x)?;
        for b in &self.l_blocks {
            y = b.forward(&y, mask)?;
        }
        Ok(y)
    }

    fn h_step(&self, z_h: &Tensor, z_l: &Tensor, mask: Option<&Tensor>) -> Result<Tensor> {
        let mut y = (z_h + z_l)?;
        for b in &self.h_blocks {
            y = b.forward(&y, mask)?;
        }
        Ok(y)
    }

    /// Refine `reps [B,T,h]` → `[B,T,h]`. The `N_cycles` of refinement run detached (the reference's
    /// `no_grad`), and only the final L+H step carries gradient — the one-step gradient approximation that
    /// makes recurrent depth trainable at this size.
    pub fn forward(&self, reps: &Tensor, mask: Option<&Tensor>, grad_last_step: bool) -> Result<Tensor> {
        let (b, t, _h) = reps.dims3()?;
        let x = self.inj.forward(reps)?;
        let mut z_l = self.z_l0.reshape((1, 1, ()))?.broadcast_as((b, t, self.z_l0.dim(0)?))?.contiguous()?;
        let mut z_h = self.z_h0.reshape((1, 1, ()))?.broadcast_as((b, t, self.z_h0.dim(0)?))?.contiguous()?;

        let cycles = if grad_last_step { self.n_cycles } else { self.n_cycles };
        for _ in 0..cycles {
            for _ in 0..self.t_inner {
                z_l = self.l_step(&z_l, &z_h, &x, mask)?;
                if grad_last_step {
                    z_l = z_l.detach();
                }
            }
            z_h = self.h_step(&z_h, &z_l, mask)?;
            if grad_last_step {
                z_h = z_h.detach();
            }
        }
        // final step (gradient flows here)
        let z_l = self.l_step(&z_l, &z_h, &x, mask)?;
        let z_h = self.h_step(&z_h, &z_l, mask)?;
        z_h + reps
    }
}


// ── embeddings-only backbone + independent heads ────────────────────────────────────────────────

/// BERT's embedding table on its own (word + position + token-type, then LayerNorm) — the reference uses
/// `BertModel.from_pretrained(base).embeddings` and discards the encoder entirely. candle keeps its
/// `BertEmbeddings` private, so this is implemented directly, which also lets it load the checkpoint's
/// `bert.embeddings.*` tensors verbatim.
pub struct BertEmbeddingsOnly {
    word: candle_nn::Embedding,
    position: candle_nn::Embedding,
    token_type: candle_nn::Embedding,
    norm: candle_nn::LayerNorm,
    pub hidden: usize,
}

impl BertEmbeddingsOnly {
    pub fn load(vb: VarBuilder, vocab: usize, max_pos: usize, type_vocab: usize, hidden: usize) -> Result<Self> {
        Ok(BertEmbeddingsOnly {
            word: candle_nn::embedding(vocab, hidden, vb.pp("word_embeddings"))?,
            position: candle_nn::embedding(max_pos, hidden, vb.pp("position_embeddings"))?,
            token_type: candle_nn::embedding(type_vocab, hidden, vb.pp("token_type_embeddings"))?,
            norm: candle_nn::layer_norm(hidden, 1e-12, vb.pp("LayerNorm"))?,
            hidden,
        })
    }

    /// `ids [B,T]` → `[B,T,h]`.
    pub fn forward(&self, ids: &Tensor) -> Result<Tensor> {
        let (b, t) = ids.dims2()?;
        let w = self.word.forward(ids)?;
        let pos_ids = Tensor::arange(0u32, t as u32, ids.device())?.reshape((1, t))?.broadcast_as((b, t))?.contiguous()?;
        let p = self.position.forward(&pos_ids)?;
        let tt = self.token_type.forward(&ids.zeros_like()?)?;
        self.norm.forward(&((w + p)? + tt)?)
    }
}

/// The reference tagger: embeddings → HRM refinement → **independent** linear heads. Each head is one
/// matrix over the shared refined state, so adding a dimension of the Vocabulary Space costs a matrix
/// rather than another encoder.
pub struct HrmTagger {
    emb: BertEmbeddingsOnly,
    core: HrmCore,
    head_a: Linear,
    head_b: Linear,
    hidden: usize,
}

/// Shape/behaviour knobs, defaulting to the reference's `cycles=2, t_inner=2, layers=2, heads=4`.
#[derive(Debug, Clone, Copy)]
pub struct HrmConfig {
    pub vocab: usize,
    pub max_pos: usize,
    pub type_vocab: usize,
    pub hidden: usize,
    pub layers: usize,
    pub heads: usize,
    pub t_inner: usize,
    pub n_cycles: usize,
}

impl HrmConfig {
    /// Defaults matching the reference `HRMTagger` on a bert-tiny embedding table.
    pub fn bert_tiny(vocab: usize, hidden: usize, max_pos: usize, type_vocab: usize) -> Self {
        HrmConfig { vocab, max_pos, type_vocab, hidden, layers: 2, heads: 4, t_inner: 2, n_cycles: 2 }
    }
}

impl HrmTagger {
    pub fn new(vb: VarBuilder, cfg: &HrmConfig, n_a: usize, n_b: usize) -> Result<Self> {
        let emb = BertEmbeddingsOnly::load(vb.pp("bert").pp("embeddings"), cfg.vocab, cfg.max_pos, cfg.type_vocab, cfg.hidden)?;
        let core = HrmCore::new(cfg.hidden, cfg.layers, cfg.layers, cfg.t_inner, cfg.n_cycles, cfg.heads, vb.pp("core"))?;
        Ok(HrmTagger {
            emb,
            core,
            head_a: candle_nn::linear(cfg.hidden, n_a, vb.pp("head_a"))?,
            head_b: candle_nn::linear(cfg.hidden, n_b, vb.pp("head_b"))?,
            hidden: cfg.hidden,
        })
    }

    pub fn hidden_size(&self) -> usize {
        self.hidden
    }

    /// Refined per-token state `[B,T,h]` — what the independent heads read, and what Head C pools spans from.
    pub fn hidden(&self, ids: &Tensor, attn: &Tensor, grad_last_step: bool) -> Result<Tensor> {
        let reps = self.emb.forward(ids)?;
        self.core.forward(&reps, Some(attn), grad_last_step)
    }

    /// `(logits_a [B,T,n_a], logits_b [B,T,n_b])`.
    pub fn forward(&self, ids: &Tensor, attn: &Tensor, grad_last_step: bool) -> Result<(Tensor, Tensor)> {
        let y = self.hidden(ids, attn, grad_last_step)?;
        Ok((self.head_a.forward(&y)?, self.head_b.forward(&y)?))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_nn::VarMap;

    fn core(h: usize, cycles: usize, t_inner: usize) -> (HrmCore, Device) {
        let device = Device::Cpu;
        let varmap = VarMap::new();
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
        let c = HrmCore::new(h, 1, 1, t_inner, cycles, 4, vb).unwrap();
        (c, device)
    }

    /// Parameter budget: embeddings dominate and the HRM core is small — the point of the architecture.
    #[test]
    fn independent_heads_and_a_small_core() {
        let device = Device::Cpu;
        let varmap = VarMap::new();
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
        let cfg = HrmConfig::bert_tiny(30522, 128, 512, 2);
        let t = HrmTagger::new(vb, &cfg, 19, 4).unwrap();
        let total: usize = varmap.all_vars().iter().map(|v| v.as_tensor().elem_count()).sum();
        let emb_params = 30522 * 128 + 512 * 128 + 2 * 128 + 2 * 128;
        eprintln!("total params {total} (embeddings {emb_params}, core+heads {})", total - emb_params);
        assert!(total < 5_000_000, "must stay a ~4M model, got {total}");
        assert!(emb_params * 10 / 8 > total - emb_params, "embeddings should dominate the budget");

        let ids = Tensor::from_vec(vec![101u32, 2054, 2003, 102], (1, 4), &device).unwrap();
        let attn = Tensor::from_vec(vec![1u32, 1, 1, 1], (1, 4), &device).unwrap();
        let (a, b) = t.forward(&ids, &attn, false).unwrap();
        assert_eq!(a.dims(), &[1, 4, 19]);
        assert_eq!(b.dims(), &[1, 4, 4]);
        // heads are independent: they read the same state but produce different widths and values
        let hidden = t.hidden(&ids, &attn, false).unwrap();
        assert_eq!(hidden.dims(), &[1, 4, 128]);
        let av: Vec<f32> = a.flatten_all().unwrap().to_vec1().unwrap();
        assert!(av.iter().all(|v| v.is_finite()));
    }

    #[test]
    fn rmsnorm_matches_the_reference_formula() {
        let device = Device::Cpu;
        let varmap = VarMap::new();
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
        let n = RmsNorm::new(4, vb.pp("n")).unwrap();
        // freshly-initialised `w` is ones in the reference; candle's default init differs, so compare the
        // normalisation itself: a scaled input must normalise to the same direction.
        let x = Tensor::from_vec(vec![1f32, 2., 3., 4.], (1, 1, 4), &device).unwrap();
        let y1 = n.forward(&x).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap();
        let y2 = n.forward(&(x.affine(10.0, 0.0).unwrap())).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap();
        for (a, b) in y1.iter().zip(y2.iter()) {
            assert!((a - b).abs() < 1e-4, "RMSNorm must be scale-invariant: {y1:?} vs {y2:?}");
        }
    }

    #[test]
    fn core_preserves_shape_and_is_residual() {
        let (c, device) = core(8, 2, 2);
        let reps = Tensor::rand(0f32, 1f32, (2, 5, 8), &device).unwrap();
        let mask = Tensor::from_vec(vec![1u32, 1, 1, 0, 0, 1, 1, 1, 1, 0], (2, 5), &device).unwrap();
        let out = c.forward(&reps, Some(&mask), false).unwrap();
        assert_eq!(out.dims(), &[2, 5, 8]);
        // output is reps + refinement, so it must differ from reps but stay finite
        let d: Vec<f32> = (out - &reps).unwrap().flatten_all().unwrap().to_vec1().unwrap();
        assert!(d.iter().all(|v| v.is_finite()), "refinement must be finite (padding mask must not produce NaN)");
        assert!(d.iter().any(|v| v.abs() > 1e-6), "core must actually change the representation");
    }

    #[test]
    fn more_cycles_change_the_refinement() {
        // depth here comes from recurrence, so cycle count must matter
        let device = Device::Cpu;
        let varmap = VarMap::new();
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
        let shallow = HrmCore::new(8, 1, 1, 1, 1, 4, vb.pp("c")).unwrap();
        let deep = HrmCore::new(8, 1, 1, 2, 3, 4, vb.pp("c")).unwrap(); // same weights, more passes
        let reps = Tensor::rand(0f32, 1f32, (1, 4, 8), &device).unwrap();
        let a: Vec<f32> = shallow.forward(&reps, None, false).unwrap().flatten_all().unwrap().to_vec1().unwrap();
        let b: Vec<f32> = deep.forward(&reps, None, false).unwrap().flatten_all().unwrap().to_vec1().unwrap();
        assert!(a.iter().zip(b.iter()).any(|(x, y)| (x - y).abs() > 1e-5), "cycle count must affect the output");
    }

    #[test]
    fn padding_mask_blocks_attention_to_pad_positions() {
        let (c, device) = core(8, 1, 1);
        // two batches: identical real prefix, different padding content
        let mut a = vec![0f32; 3 * 8];
        for (i, v) in a.iter_mut().enumerate() {
            *v = (i % 7) as f32 * 0.1;
        }
        let reps_a = Tensor::from_vec(a.clone(), (1, 3, 8), &device).unwrap();
        // same first two positions, wildly different third (which is masked out)
        let mut b = a.clone();
        for v in b[16..24].iter_mut() {
            *v = 99.0;
        }
        let reps_b = Tensor::from_vec(b, (1, 3, 8), &device).unwrap();
        let mask = Tensor::from_vec(vec![1u32, 1, 0], (1, 3), &device).unwrap();
        let oa = c.forward(&reps_a, Some(&mask), false).unwrap().narrow(1, 0, 2).unwrap();
        let ob = c.forward(&reps_b, Some(&mask), false).unwrap().narrow(1, 0, 2).unwrap();
        let va: Vec<f32> = oa.flatten_all().unwrap().to_vec1().unwrap();
        let vb2: Vec<f32> = ob.flatten_all().unwrap().to_vec1().unwrap();
        for (x, y) in va.iter().zip(vb2.iter()) {
            assert!((x - y).abs() < 1e-3, "masked positions must not influence real ones: {va:?} vs {vb2:?}");
        }
    }
}