Module hrm 

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.

Structs

BertEmbeddingsOnly

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.

HrmConfig

Shape/behaviour knobs, defaulting to the reference’s cycles=2, t_inner=2, layers=2, heads=4.

HrmCore

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.

HrmTagger

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.

RmsNorm

x * rsqrt(mean(x²) + eps) * w — the reference’s RMSNorm (no mean subtraction, learned gain).