Brujula-450M-Retrieval

Brújula-450M-Retrieval

A long-context needle-retrieval fine-tune of Brújula-450M: find a fact buried in tens of thousands of tokens of narrative prose. It takes the base from 0% to 93% on BABILong qa1 at 16K, and holds ~80% out to 64K — an 8× extrapolation beyond the 8K window it was trained at.

⚠️ Before you download this: it fails on real documents

That 93% is real, and it is also misleading on its own. This model is a synthetic-needle specialist. Fed an actual research paper as a 24K-token haystack with 8 factual questions, it scores 0/8 — and answers most of them with a word like "garden", which is a bAbI location, not anything in the document.

• Want to ask questions about a real paper or document? Use Brújula-450M-DocQA instead. This is the wrong model for that, and no prompt will fix it.
• Want a long-context needle benchmark to work? This does that, well, at up to 64K.

The gap between "93% on BABILong" and "0/8 on one real paper" is the most useful thing here, and it is why this model is published at all rather than quietly kept. Details in The failure that matters.

Results

BABILong — bAbI facts embedded in real PG-19 book prose. Accuracy %, qa1 / qa2 / qa3:

The base model scores 0 / 0 / 0 at every one of these lengths. So does a version fine-tuned on naive single-template passkey needles — see below.

Trained at block 8192; 16K/32K/64K are all extrapolation via YaRN. "Train small, test big" holds here, and reach scales with model size: the same recipe on the 150M gets 73% at 16K and fades to ~50% by 32K.

The failure that matters

Fed the real DeepSeek-V2 paper as a 24K-token haystack with 8 factual questions, this model scores 0/8 — and answers most of them with a word like "garden", which is a bAbI location. It is not reading the prose. It is pattern-matching the User:/Question:/Assistant: frame to its synthetic training distribution and emitting a canned answer.

BABILong is bAbI structure wrapped in book prose, which is close enough to the training distribution to work. A technical paper is genuinely out of distribution, and the model has nothing to fall back on.

This is worth publishing precisely because the BABILong numbers look convincing on their own. A benchmark score of 93% at 16K coexists with 0/8 on the first real document we tried. Synthetic long-context benchmarks can overstate real long-context ability by that much.

A second, related result: an earlier fine-tune of this base on single-template passkey needles ("the X is Y → copy Y") reached 87/93/83% on passkey at 16/32/64K while scoring 0% on BABILong. That taught template-copying, not search. Fixing it needed dictionary-search data that forces an actual lookup, plus real narrative training data — which is the recipe below.

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "Sakatepon/Brujula-450M-Retrieval"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    repo, trust_remote_code=True, dtype=torch.bfloat16
).eval()

haystack = "..."      # long context
question = "Where is the apple?"

prompt = f"User: {haystack.strip()}\n\nQuestion: {question.strip()}\nAssistant: "
ids = tok(prompt, return_tensors="pt")
out = model.generate(**ids, max_new_tokens=24, do_sample=False, use_cache=False,
                     eos_token_id=50256, pad_token_id=50256)
print(tok.decode(out[0][ids.input_ids.shape[1]:], skip_special_tokens=True))

Context regime. config.json ships yarn_ms with rope_trained_len=1024 and rope_scale_len=34816, which is the regime the 16K and 32K numbers were measured in. For 64K, rebuild the scaling for a table that covers it — load the config, set rope_scale_len=68000, and pass it to from_pretrained. Using a per-length table instead of one fixed table changes the RoPE angles and will cost you accuracy.

use_cache=False is required (Block Attention-Residuals mix across layers). At long context, keep batch size 1 and do not materialize full-vocabulary logits.

Training

LoRA rank 16 on attention + FFN projections, base frozen, adapters merged for release. fp32 (TF32), block 8192, answer-masked loss, yarn_ms RoPE scaling at scale_len 32768 from a pre-train length of 1024. Rented GPU, ~$6. Data was a mix of:

• procedurally generated dictionary search — 85+ integer dictionaries, asked for one key's value and location, which forces a genuine lookup rather than a copy
• BABILong-train qa1/qa2/qa3 — real narrative prose with latest-location state tracking

at mixed lengths. Both halves were necessary: dictionary-search alone teaches copying.

Limitations

• Fails on real documents (0/8 on a real paper, see above). This is the headline caveat.
• Emits bAbI-flavoured answers (locations, objects) when out of distribution.
• qa2/qa3 (multi-fact, state-tracking) lag qa1 substantially — 60-ish vs 93.
• No KV cache in this export; long-context generation is slow.
• English only, GPT-2 BPE.

License

Apache-2.0.