epfl-dlab / epfl-dlab/zip2zip

Eval adapter: LZW boundary corruption + codebook state leak + logit masking

Open
#15 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
19
Forks
6
PR merge metrics
No merged PRs in 30d

Description

Found while auditing zip2zip-core's eval paths against the pip package. HIGH → LOW.

### HIGH

- [ ] **MC scoring can drop continuation tokens at the context boundary (LZW isn't
prefix-stable).** `Zip2ZipForLMEval` uses lm-eval's stock `_encode_pair`:
`encode(ctx+cont)` sliced at `len(encode(ctx))`. LZW's pending buffer at the ctx end
can merge with the first continuation tokens, so those base tokens go missing from
**both** the model input and the scored span. Minimal repro (Phi config,
`initial_vocab_size=32011`):
```
ctx = [10,20,30,10,20,30,5,10] -> ctx_enc = [10,20,30,32011,30,5,10] (7 toks)
ctx+cont([20,30,7]) -> whole_enc = [10,20,30,32011,30,5,32014,7]
whole_enc[:7] != ctx_enc; continuation_enc = [7] -> base tokens 20,30 vanish
```
Affects every MC number from this adapter (small, data-dependent; likely shared with
the original eval-fork lineage, so existing baselines stay self-consistent). Rolling
perplexity has the same issue at each window boundary
(`harness.py:186` `loglikelihood_rolling`, boundary at `215`). *Fix:*
override `_encode_pair`/`_loglikelihood_tokens` — LZW-encode the full text once and
locate the continuation by base-token span (as zip2zip-core does). **Note: this
changes comparability with published numbers — coordinate before re-baselining.**
(`harness.py:170` `_loglikelihood_tokens` delegates to the stock HFLM path;
root cause: end-of-input flush in `zip2zip-compression/src/codec.rs` ~line 280)

- [ ] **Codebook state leaks into the first `generate()` after a scoring phase.**
`generate` calls `init_codebooks_and_hyper_weight_cache` (a no-op when caches
exist); reset happens only AFTER generate, and `base_model.generate` bypasses the
wrapper-forward reset. So the first generation after the loglikelihood phase inherits
the Rust LZW state + stale hyper caches. Impact: 1/1319 GSM8K samples in a full run,
**1/20 (5%) in `--limit 20` smoke runs** — enough to flip smoke comparisons. *Fix:*
call `codebook_manager.reset()` at the top of `generate`. (Related: the
`generate_until` fallback re-decodes the suffix with a fresh LZW state, mis-expanding
prompt-referencing hyper ids — `harness.py:90-94`.) (`model.py:321,325,329`
generate/init/reset-after; `codebook.py:51-55` init is a no-op when caches exist)

### MEDIUM

- [ ] **Unused codebook slots contribute logit 0 to every softmax.** `HyperLinear` has
no used-entry mask; zero-init unused slots each add `exp(0)` to the denominator (up
to `max_codebook_size` spurious terms). zip2zip-core masks them to −inf. Scores of
core-trained exports are systematically deflated and argmax can pick an unused slot;
pip-trained models (released v0.1) are self-consistent. Means pip and core numbers
aren't 1:1 comparable. *Fix:* track used entries and mask unused hyper logits to
−inf (behind a flag to preserve parity with old-convention models).
(`nn/linear.py:31-47`, `codebook.py:98-126`)

- [ ] **`from_pretrained` falls back silently on load failures.** (1) hyper-encoder
load wrapped in bare `except Exception` → INFO "No hyper encoders found" and
continues with RANDOM encoders; (2) any PEFT `OSError/FileNotFoundError/ValueError`
→ treated as "no adapter", scores the BASE model; (3) decoder loads `strict=False`,
warnings only. All produce plausible-looking wrong numbers under file/version skew.
*Fix:* narrow the excepts to genuine not-found; raise on strict-load failures;
escalate missing decoder-layer keys to errors. (`model.py:204`, `392-397`, `413-421`)

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the Phi minimal repro and a --limit 20 GSM8K smoke run, then read harness.py, model.py, codebook.py, nn/linear.py, and the LZW flush logic in zip2zip-compression/src/codec.rs. Verify the boundary, generation-state, masking, and load-failure behaviors before changing them. Done means corrected evaluations and loading errors without silently producing plausible wrong scores, with results checked against the stated parity and baseline caveats.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, rust
Domain
machine-learning, testing-qa
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.