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:
| context | qa1 | qa2 | qa3 |
|---|
| 16K | 93 | 63 | 60 |
| 32K | 80 | 57 | 73 |
| 64K | 80 | 60 | 63 |
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
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4repo = "Sakatepon/Brujula-450M-Retrieval"
5tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
6model = AutoModelForCausalLM.from_pretrained(
7 repo, trust_remote_code=True, dtype=torch.bfloat16
8).eval()
9
10haystack = "..." # long context
11question = "Where is the apple?"
12
13prompt = f"User: {haystack.strip()}\n\nQuestion: {question.strip()}\nAssistant: "
14ids = tok(prompt, return_tensors="pt")
15out = model.generate(**ids, max_new_tokens=24, do_sample=False, use_cache=False,
16 eos_token_id=50256, pad_token_id=50256)
17print(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.