HRM-Embed

A text-embedding model built on a reasoning architecture.

🤗 Model on Hugging Face · RakanLabs

An embedding model built on top of the HRM architecture by Sapient, following the same recipe as RakanEmbed4B.¹

• 🧠 Unusual backbone, standard recipe: a depth-recurrent HRM instead of a vanilla transformer encoder.
• 🔬 Reproducible eval: the full BRIGHT harness is included (raw / rewrite / merged, per domain).
• 📊 Benchmarked on BRIGHT: one of the hardest reasoning-retrieval benchmarks out there.

The hypothesis

BRIGHT is a reasoning-retrieval benchmark: matching a query to the right document often takes several hops of reasoning, not just surface similarity. The bet was simple, that a depth-recurrent architecture, one that loops over its own hidden state instead of doing a single forward pass, is a natural fit for that kind of multi-hop reasoning. That part held up: the looping helps on the reasoning-heavy side.

The limit turned out to be knowledge, not reasoning. The base model was pretrained on a relatively narrow slice of text, so it simply hasn't seen enough of the world to retrieve well on knowledge-heavy domains.

How it works

Nothing here is a new technique; it is an amalgamation of standard ones on an unusual backbone. The embedding recipe is deliberately standard: mean-pool the final hidden state and L2-normalize (the Sentence-BERT convention), run the model bidirectionally (the causal-to-encoder conversion popularized by LLM2Vec), and train it contrastively (InfoNCE). The only unusual part is the backbone: applying that recipe to a depth-recurrent HRM and pooling the recurrence state z_h.

Install

pip install -r requirements.txt

Quickstart

python examples/embed.py

The model is not a sentence-transformers model. It loads via trust_remote_code, and you pool the recurrence state yourself. Minimal usage:

import torch, torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM

name = "viventhraa96/HRM-Embed-0.6b"          # HF repo id or a local path
dev = "cuda" if torch.cuda.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained(name)
if tok.pad_token is None: tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(
    name, trust_remote_code=True, torch_dtype=torch.bfloat16).to(dev).eval()
model.lm_head = torch.nn.Identity()        # embeddings come from z_h, not the LM head

@torch.no_grad()
def embed(texts, max_length=512):
    tok.padding_side = "right"
    e = tok(texts, truncation=True, max_length=max_length, padding=True, return_tensors="pt").to(dev)
    pos = torch.arange(e.input_ids.shape[1], device=dev).unsqueeze(0).expand(e.input_ids.shape[0], -1)
    z, _ = model.model(e.input_ids, position_ids=pos, use_cache=False, token_type_ids=e.attention_mask)
    m = e.attention_mask.unsqueeze(-1).to(z.dtype)
    return F.normalize(((z * m).sum(1) / m.sum(1).clamp_min(1)).float(), p=2, dim=-1)  # [N, 1280]

See examples/embed.py for a runnable version.

Benchmarks: BRIGHT (reasoning retrieval)

Mean nDCG@10 over BRIGHT's 12 domains, for three query modes: raw (bare embedder), rewrite (an LLM rewrites the query first, as most top BRIGHT systems do)², merged (raw + rewrite).

Query mode	Mean nDCG@10
raw (bare embedder)	18.1
+ query rewriting	34.3
merged (raw + rewrite)	33.7

Per-domain numbers are in evals/bright/ (results.json and <domain>/results.json). Query rewriting lifts every domain except code retrieval (LeetCode: raw 22.6 to rewrite 12.9, since expanding a terse problem spec into prose moves the query off the corpus distribution; use merged for code).

Strong at: theorem/definition/reference lookup (TheoremQA), vocabulary-aligned scientific QA (biology, psychology). Weak at: reasoning-transfer retrieval (aops), community/procedural QA (robotics, stackoverflow). pony is the extreme case: near-chance without rewriting (raw 1.1) yet among the strongest with it (46.5), the most rewrite-dependent domain in the set. This is a small model on a deliberately adversarial benchmark; the absolute scores are modest and reported plainly.

Repository layout

examples/embed.py                     # minimal inference example
evals/bright/                         # BRIGHT eval harness
  eval_bright_hrm_lancedb.py           # LanceDB eval, raw/rewrite/merged
  run_one_domain.sh                    # run one domain and refresh the board
  results.json (combined board), <domain>/results.json (per-domain)

Reproduce

• Eval (raw variant reproduces from public BRIGHT alone):

  python evals/bright/eval_bright_hrm_lancedb.py --checkpoint /path/to/model --domains biology

  See evals/bright/README.md for the rewrite/merged variants.

Credits

The Hierarchical Reasoning Model (HRM) architecture is by Sapient Intelligence (sapientinc/HRM, sapientinc/HRM-Text, arXiv:2506.21734). All architectural credit is theirs. This model was fine-tuned from an open, Apache-2.0 Xiaoye08/HRM-Text-0.6B pretrained checkpoint; the HRM-Text pretraining pipeline is described in arXiv:2605.20613. Please cite the original HRM if you build on this.

License

Apache-2.0. This is a derivative of Apache-2.0-licensed weights; attribution to Sapient Intelligence (HRM) and the HRM-Text project is preserved above.

Footnotes

1. The entire model, including weights, training, and evaluation, was produced on a single consumer GPU (an 8 GB NVIDIA RTX 3060 Ti). ↩

2. Rewritten queries for the rewrite and merged variants come from INF-X-Retriever's rewrite_data. ↩