The full-recompute (FR) upper bound is understated by three prompt-construction issues in the demo harness

[cntlgns]

TL;DR. Running the shipped e2e harness, the fr (full-recompute) accuracy
upper bound on the QA tasks comes out surprisingly low, which makes the
headline "LegoLink recovers ≈all of FR" look better than it is. The cause appears
to be three independent prompt-construction details in the demo harness — not
the LegoLink algorithm itself:

1. a begin-of-text marker is repeated in the middle of the sequence (one per chunk),
2. the assistant turn is never opened, so chat-template text leaks into the scored output, and
3. the question is wrapped in a redundant, length-capped instruction template.

Fixing all three raises fr sharply (e.g. hotpotqa fr 0.324 → 0.541) while
LegoLink rises much less, exposing a previously-hidden recovery gap: best
kvlink-k / fr ≈ 78% (2wikimqa), 91% (hotpotqa), 90% (musique) rather than
"near the upper bound". The same effect reproduces on the repo's default model
(Mistral-7B). Is the multi-begin-of-text / no-assistant-opener prompt
construction intended, or an artifact of the demo harness?

Scope and setup

All numbers below are on the LongBench QA subsets shipped in this repo (200
curated samples per task — referred to here as LB200): 2wikimqa,
hotpotqa, musique, needle. Model: meta-llama/Llama-3.1-8B-Instruct.
Metric: the repo's own token-level F1 (qa_f1_score). Decoding: greedy,
max_tokens=32. Approaches:

• fr — full recompute (every token attends to every other token); this is
  the intended accuracy upper bound.
• naive — reuse cached KV with no re-encoding.
• kvlink-k — LegoLink with k re-encoded initial tokens per chunk
  (k = 2, 4, 8, 16, 32).

The point of this report: the released harness builds the prompt in a way that
depresses the fr upper bound itself, so the headline "LegoLink recovers
~all of FR" looks better than it is. Three independent issues are responsible.
Fixing them raises fr substantially while LegoLink rises much less, revealing a
real recovery gap.

How to reproduce

Stock harness, no modifications. Run one fresh process per (dataset, approach) —
kvlink/cacheblend mutate the cached KV in place, so approaches cannot be reused
within a single process:

python3 benchmarks_ours/evals/e2e/main.py \
  --dataset hotpotqa --model meta-llama/Llama-3.1-8B-Instruct \
  --approach fr --tot_num_data 200
# repeat with --approach  naive | kvlink-2 | kvlink-4 | kvlink-8 | kvlink-16 | kvlink-32
# and with  --dataset  2wikimqa | hotpotqa | musique | needle

The "after fix" tables (§3) are the same commands after applying the three
edits described in §2. (§5 uses the repo's default mistralai/Mistral-7B-Instruct-v0.2.)

1. Performance of the original implementation

benchmarks_ours/evals/e2e/main.py as published, no modifications, F1:

benchmark	naive	fr	kvlink-2	kvlink-4	kvlink-8	kvlink-16	kvlink-32
2wikimqa	0.054	0.221	0.211 (95%)	0.215 (97%)	0.216 (98%)	0.194 (88%)	0.200 (90%)
hotpotqa	0.045	0.324	0.256 (79%)	0.247 (76%)	0.261 (81%)	0.243 (75%)	0.282 (87%)
musique	0.026	0.187	0.140 (75%)	0.153 (82%)	0.153 (82%)	0.151 (81%)	0.145 (78%)
needle	0.015	0.677	0.599 (88%)	0.601 (89%)	0.616 (91%)	0.639 (94%)	0.664 (98%)

Note how close the best kvlink-k is to fr on the QA tasks (e.g. hotpotqa
fr 0.324 vs kvlink-32 0.282). That apparent closeness is largely an artifact
of an fr baseline that is itself too low.

Same harness on the full source datasets (full2wiki / fullhotpot / fullmusique — complete question sets)

benchmark	naive	fr	kvlink-2	kvlink-4	kvlink-8	kvlink-16	kvlink-32
full2wiki	0.090	0.248	0.262 (106%)	0.261 (106%)	0.262 (106%)	0.218 (88%)	n/a
fullhotpot	0.099	0.368	0.301 (82%)	0.301 (82%)	0.302 (82%)	0.271 (74%)	0.288 (78%)
fullmusique	0.021	0.157	0.108 (69%)	0.110 (70%)	0.111 (71%)	0.082 (52%)	0.096 (61%)

(full2wiki kvlink-32 is omitted.) On full2wiki the best kvlink-k even
exceeds the original fr — a direct symptom of an fr baseline depressed by
the issues in §2, not of LegoLink being genuinely better than full recompute.

2. Three issues in the prompt construction

2.1 Several "begin-of-text" tokens are placed in the middle of the sequence

The harness tokenizes the system prompt, each document, and the question as
separate chunks. The tokenizer adds a leading begin-of-text marker
(tokenizer.bos_token_id, <|begin_of_text|>) to every chunk. When the chunks
are concatenated into the final sequence, the code keeps every chunk's
leading marker instead of keeping a single one at the very front:

• main.py:149 — the line actually used:

  input_ids = [_token_ids[i] for _token_ids in token_ids for i in range(len(_token_ids))]

  range(len(...)) starts at index 0, so each chunk contributes its own
  begin-of-text. The standard single-marker version is present but
  commented out right above at main.py:146.
• The same thing happens when the KV cache is collected: main.py:177-178
  keep each chunk's leading-marker key/value (the line that would drop it is
  commented out at main.py:172-173).

A begin-of-text token appearing repeatedly inside the context is
out-of-distribution for the model and degrades the full-attention pass.

Fix: keep only one begin-of-text at position 0; drop the per-chunk leading
marker from both the input ids and the collected KV cache so the sequence is a
single, standard prompt.

2.2 The model is never told to start answering (no assistant-turn opener)

After the question, the harness appends a single end-of-turn token and stops:

• main.py:142:

  free_form_token_ids = tokenizer.encode(free_form_prompt) + [tokenizer.eos_token_id]

For a Llama-3 chat model, a turn is opened with a header such as
<|start_header_id|>assistant<|end_header_id|>. Because the harness never emits
the assistant-turn opener
(<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n), the model is
left to generate that header text itself before it produces the answer. That
header/role text then ends up inside the decoded string, which is scored
verbatim — the output is taken with no cleanup at main.py:250:

results[f'output_{configs.approach}'].append(output[0].outputs[0].text)

The leaked template text lowers token-F1 (and on some tasks collapses it).

Fix: close the user turn and open the assistant turn by appending
<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n before generating,
and strip any residual template markers from the decoded text before scoring.

2.3 The question is wrapped in a redundant, answer-shaping instruction

Each question is not passed as-is; it is wrapped in a fixed instruction
template. For example:

• benchmarks_ours/data_sets/wikimqa.py:74 (identical at
  hotpotqa.py:25, musique.py:26; needle.py:161 uses "20 words"):

  row['input'] = '\n\nAnswer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: ' + row['input'] + '\nAnswer within 5 words:'

Two problems: (a) this instruction duplicates the one already placed in the
system prompt (*_append_system_prompt, e.g. wikimqa.py:78), and (b) the
trailing Answer within 5 words: constraint pushes the model to reword or pad
its answer to fit, which hurts token-level F1 against the gold answer.

Fix: feed the bare question, with no surrounding instruction template.

A related fourth detail: the system prompt is supplied as a user-role
message (main.py:137, apply_chat_template([{"role": "user", ...}])), so it
lands inside the user turn rather than a real system turn. We also tried
moving/shortening the system prompt; in controlled runs this changed F1 by
≤0.001, so it is not a driver and is omitted from the fixes above.

3. Performance after the three fixes

Same benchmarks, same model, same metric, same max_tokens, with all three
fixes in §2 applied:

benchmark	naive	fr	kvlink-2	kvlink-4	kvlink-8	kvlink-16	kvlink-32
2wikimqa	0.042	0.450	0.254 (56%)	0.297 (66%)	0.349 (78%)	0.317 (70%)	0.342 (76%)
hotpotqa	0.091	0.541	0.449 (83%)	0.461 (85%)	0.454 (84%)	0.457 (84%)	0.494 (91%)
musique	0.044	0.298	0.228 (77%)	0.235 (79%)	0.269 (90%)	0.239 (80%)	0.258 (87%)
needle	0.226	0.712	0.821 (115%)	0.814 (114%)	0.816 (115%)	0.809 (114%)	0.790 (111%)

Full-source benchmarks (full2wiki / fullhotpot / fullmusique — confirms the effect is not a small-slice artifact)

The LB200 sets above are 200 curated samples per task. To confirm the effect is
not an artifact of that small, curated slice, we re-ran the QA tasks on the
full source datasets (the complete question sets, not the 200-sample
LongBench cut), same model and metric, with all three fixes applied:

benchmark	naive	fr	kvlink-2	kvlink-4	kvlink-8	kvlink-16	kvlink-32
full2wiki	0.062	0.504	0.400 (79%)	0.402 (80%)	0.423 (84%)	0.347 (69%)	n/a
fullhotpot	0.066	0.682	0.470 (69%)	0.471 (69%)	0.494 (72%)	0.468 (69%)	0.518 (76%)
fullmusique	0.017	0.361	0.182 (50%)	0.214 (59%)	0.218 (60%)	0.192 (53%)	0.235 (65%)

(full2wiki kvlink-32 is omitted.) The recovery gap is, if anything, wider on
the full sets: best-kvlink-k / fr is 84% (full2wiki), 76% (fullhotpot),
65% (fullmusique) — confirming the LB200 finding rather than depending on it.

4. Effect

The fr upper bound rises sharply on every QA task — 2wikimqa 0.221 → 0.450,
hotpotqa 0.324 → 0.541, musique 0.187 → 0.298 — while needle (already
saturated) is unchanged. Crucially, LegoLink rises far less than fr, so the
best-kvlink-k / fr recovery ratio drops from "near the upper bound" to
roughly 78% (2wikimqa), 91% (hotpotqa), 90% (musique) — i.e. there is a real,
previously-hidden gap between LegoLink and a faithful full-recompute baseline.
(needle is a separate regime: with a correct prompt LegoLink slightly exceeds
fr.)

In short, the three issues are independent and each is a small, local change to
the prompt construction; together they correct the fr upper bound and the
LegoLink-vs-FR comparison.

5. The same effect on the repo's default model (Mistral-7B-Instruct-v0.2)

Everything above is on Llama-3.1-8B. The repo's e2e scripts actually default to
mistralai/Mistral-7B-Instruct-v0.2 (benchmarks_ours/evals/e2e/run_all.sh,
run_one.sh), so we repeated the analysis there — same metric (qa_f1_score),
same decoding (greedy, max_tokens=32). Takeaway up front: both the
understated fr upper bound and the LegoLink-vs-FR recovery gap reproduce; two
of the three fixes (§2.1, §2.2) transfer directly, while §2.3 is a token-F1
artifact specific to how Mistral handles the length instruction.

Mistral-7B — full tables & the 2×2×2 fix ablation

5.1 Which of the three fixes transfer

We ran a full 2×2×2 ablation over the three fixes and measured the fr upper
bound on the LB200 QA tasks. Average marginal effect of each fix on fr
(token-F1, averaged over 2wikimqa/hotpotqa/musique and over the other two
flags):

fix	marginal effect on fr
§2.1 single begin-of-text	+0.031
§2.2 assistant-turn opener	+0.015
§2.3 bare question	−0.057

§2.1 and §2.2 help Mistral as well, but §2.3 (bare question) hurts here:
without the Answer within 5 words: instruction, Mistral returns
correct-but-verbose answers that token-F1 penalizes on precision (this is not
template leakage — we verified zero [INST]/header text in the outputs). The
best subset for Mistral is therefore §2.1 + §2.2, which lifts the QA-average
fr from 0.207 → 0.271 (+31%); adding §2.3 cancels the gain (back to 0.207).

§2.2 is model-specific text. Mistral's chat template uses [INST] … [/INST],
so the assistant-turn opener is the closing [/INST] rather than Llama's
<|eot_id|><|start_header_id|>assistant…. We derive it from the tokenizer's
own chat template instead of hardcoding, so the same fix is correct for both
models.

5.2 Original implementation (Mistral), LB200

benchmark	naive	fr	kvlink-2	kvlink-4	kvlink-8	kvlink-16	kvlink-32
2wikimqa	0.037	0.185	0.157	0.152	0.155	0.147	0.152
hotpotqa	0.035	0.294	0.216	0.218	0.215	0.222	0.235
musique	0.027	0.142	0.095	0.100	0.110	0.106	0.118
needle	0.179	0.584	0.420	0.435	0.421	0.440	0.438

5.3 After §2.1 + §2.2 (single begin-of-text + assistant opener; §2.3 NOT applied), LB200

benchmark	naive	fr	kvlink-2	kvlink-4	kvlink-8	kvlink-16	kvlink-32
2wikimqa	0.041	0.244	0.191	0.195	0.196	0.183	0.166
hotpotqa	0.038	0.366	0.238	0.238	0.272	0.275	0.262
musique	0.028	0.204	0.154	0.155	0.163	0.157	0.165
needle	0.316	0.715	0.725	0.709	0.715	0.717	0.716

fr rises on every QA task (2wikimqa 0.185 → 0.244, hotpotqa 0.294 → 0.366,
musique 0.142 → 0.204) and on needle (0.584 → 0.715), exactly as on Llama —
while best-kvlink-k rises far less.

5.4 The recovery gap holds, and is visible even in the unmodified harness

best-kvlink-k / fr (recovery ratio). We also ran the full standard
benchmarks (full validation sets: full2wiki 12,576, fullhotpot 7,405,
fullmusique 2,417) under the same two settings:

benchmark (data)	fr orig	recov% orig	fr fixed	recov% fixed
2wikimqa (LB200)	0.185	85	0.244	80
hotpotqa (LB200)	0.294	80	0.366	75
musique (LB200)	0.142	83	0.204	81
2wikimqa (full)	0.260	84	0.295	87
hotpotqa (full)	0.429	69	0.427	71
musique (full)	0.167	67	0.233	70

On every QA benchmark, LegoLink recovers only 70–87% of the full-recompute
upper bound — and unlike Llama, on Mistral this gap is already visible in the
unmodified harness (67–85%), independent of the fr correction. needle
is the separate regime (LegoLink ≈ fr, recovery 101%); multi_news (Rouge-L)
is flat (≈99%).

Two of the three prompt fixes transfer directly to Mistral; the third (§2.3) is
a token-F1 artifact specific to how Mistral handles the length instruction, so
for Mistral we report the two-fix subset (§2.1 + §2.2).