Skip to content

Evaluation windows sliced mid-document lack BOS; accounts for the reported perplexity reduction #1

Description

@JackCai1206

Summary

eval.py::sample_windows slices evaluation windows out of an already-tokenised
document. Windows starting at index > 0 therefore carry no BOS token. With
the documented --positions 3, two of every three windows are affected, and
those windows carry essentially the entire reported perplexity reduction.

Adding BOS back to sliced windows drops the baseline from 40.93 → 21.30 and
reduces the α=0.07 / 12-5 effect from −13.49% → +0.27%.

Where

# eval.py, sample_windows()
ids = tokenizer(text, truncation=True, max_length=window*8,
                return_tensors="pt").input_ids[0]   # BOS is at index 0
starts = [0]                                        # this window HAS BOS
for _ in range(positions - 1):
    starts.append(rng.randint(1, max_start))        # these do NOT
for s in starts:
    windows.append(ids[s : s + window].to(device))

Gemma-3 is trained with BOS at position 0 of every sequence. Without it the
first content token is forced to serve as the attention sink, which corrupts
representations and inflates perplexity substantially.

Evidence

Running your eval.py unmodified reproduces your published numbers exactly
(results.json records baseline_ppl: 40.94; we measure 40.9336).
Splitting those 18 windows by whether they contain BOS:

group n mean baseline ppl
starts=[0] (has BOS) 6 20.32
sliced (no BOS) 12 73.87 (max 165.7)

A/B with a one-line fix, same seed, same documents, same configs:

config original BOS-fixed
baseline ppl 40.93 21.30
α=0.07, 12-5 −13.49% +0.27%
α=0.07, 11-4 −6.15% +0.30%
α=0.04, 12-5 −8.80% +0.02%
α=0.04, 11-4 −4.24% +0.05%

Internal control: the six starts=[0] windows already contained BOS, so the fix
should not change them — and it does not (10.908, 26.789, 20.096, 26.148, 18.441, 19.511 in both arms, identical). Only the sliced windows move, and they
move a lot: 165.7→48.0, 148.3→35.7, 119.1→44.0, 95.8→26.7.

Independent cross-check: the paper's Table 1 reports a Gemma3-1B / PG-19 baseline
of 22.27. The BOS-fixed baseline (21.30) matches that within 4%; the original
(40.93) is 84% off. So the paper's own evaluation appears to be correctly
tokenised, and the fixed configuration is the one consistent with it.

Suggested fix

for s in starts:
    w = ids[s : s + window]
    if s > 0 and tokenizer.bos_token_id is not None:
        w = torch.cat([torch.tensor([tokenizer.bos_token_id]), w[:-1]])
    windows.append(w.to(device))

(Or restrict windows to document starts, or re-tokenise each window's decoded
text with add_special_tokens=True.)

Scope

This concerns the reproduction harness only, not the paper. Mozer et al.
report healthy baselines (e.g. 14.76 for Gemma3-4B), inconsistent with a
BOS-less evaluation, so their result may well be correct — we have not been able
to test it, since no official code exists. We are separately attempting the
paper's Appendix B.2 protocol and will share results if useful.

Thanks for publishing this — the code is clear and well documented, which is the
only reason the issue was findable at all.

Environment

torch 2.11.0+cu128 / transformers 4.57.6 / datasets 5.0.1, NVIDIA H100 80GB,
google/gemma-3-1b-pt, PG-19 (emozilla/pg19-test), --n_docs 6 --positions 3 --window 512 --seed 0.

Two unrelated notes while running it: recirculation.py --text crashes with a
device mismatch (ids on CPU, logits on CUDA) at the perplexity_from_logits
call; and datasets 5.x no longer supports streaming=True against the
emozilla/pg19 script dataset in offline mode.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions