Fix AIM row alignment across batch and head boundaries

[ybochkov]

Summary

get_aim_states and get_aim_star_states flatten the (batch, head) structure and detect word boundaries with roll() and a global cumsum(). Adjacent rows in the flattened tensor can then leak into each other: whenever the word ids touching a row boundary (or the wrap-around between the tensor's first and last elements) happen to be equal, word pairs from unrelated rows or heads are merged into one pair row or dropped entirely. Because repeat_interleave places the per-head copies of the same row next to each other, a colliding row collides at every head boundary at once.

Teacher and student AIM states are built independently, so one merged or dropped row shifts every following row and breaks their one-to-one correspondence in the loss:

• when the teacher/student pair counts end up different, the size-mismatch guard in common_step skips the batch with a dummy loss;
• when the counts happen to match, misaligned rows are paired in the MSE/cosine loss with no error at all;
• a batch containing no completed word pair crashes on valid_word_ids.max() of an empty tensor.

Severity depends on the tokenizer. When every sequence starts with a special token mapped to -100 (Llama/Gemma/Mistral-style BOS), the leading -100 acts as a separator at each flattened row boundary, and randomized synthetic batches (right-padded, longest row unpadded, 1–4 heads) show no divergence from the fixed implementation — past runs with such models were most likely unaffected. For tokenizers without a leading special token (e.g. Qwen3, which this repo ships an adapter for), the same randomized stress test produced wrong AIM/AIM* states in roughly 10–15% of batches and crashed in 4%.

Fix

Assign pair ids per (batch row, head) slice: the new causal_word_pair_ids helper keeps the row/head structure until the very end, so boundary collisions are impossible by construction, and get_aim_states returns an empty state tensor when a batch contains no complete word pair instead of crashing. A shared last_token_mask helper replaces the roll-based last-token detection in both AIM and AIM*.

Behavior is preserved where the previous code was correct: on randomized BOS-style batches the old and new implementations agree exactly (200/200 trials).

Note: this removes the implementation-induced pair mismatches, not every possible "Word states size mismatch" — teacher and student are tokenized independently, so truncation can still produce genuinely different word sets; that remains the guard's job.

Tests

tests/test_aim_impl.py adds regression tests that keep identical word ids separate across batch rows and across attention heads, for both AIM and AIM*. All four fail on the previous implementation (one via the empty-tensor crash) and pass with the fix.

🤖 Generated with Claude Code