[Bug Report] TransformerBridge ignores attention_mask when deriving position_ids, so left-padded input silently returns wrong logits and loss

[sohv]

Describe the bug

TransformerBridge.forward() does not derive position_ids from a supplied attention_mask, so left-padded input silently gets the wrong absolute positions. No error, no NaN — just wrong logits and a wrong loss. HookedTransformer handles this correctly, so the same code returns different numbers after migrating to the bridge.

Not a duplicate of #1607 - That issue is the attention_mask never reaching loss_fn, which affects loss aggregation only but the logits are correct. This is the mask never reaching position_ids, which corrupts the logits themselves and the loss downstream. They are independent and they compound — with #1608 applied (which fixes #1607), max |logit diff| here is still 1.361e+02, and with this fix alone the loss is still wrong by 1.343308. Both are needed for correct left-padded batches; measurements in "Interaction with #1607 / #1608" below.

gpt2, one prompt, mask supplied, varying the number of left pads:

n_pad	HT loss	Bridge loss	HT drift	Bridge drift
0	4.503170	4.503170	—	—
1	4.503169	13.594296	1.4e-06	9.09
2	4.503170	11.478184	4.8e-07	6.98
3	4.503169	11.154946	9.5e-07	6.65
5	4.503170	11.113400	4.8e-07	6.61
8	4.503169	10.787075	1.4e-06	6.28

Right padding is unaffected for both (drift ≤ 9.5e-07), which is expected — causality already protects it. The failure is specific to left padding, where every real token's absolute position shifts by the pad count.

Not model-specific. distilgpt2, n_pad=3: HT drift 4.8e-07, Bridge drift 4.417137, max |logit diff| 5.791e+01.

The magnitude scales with how much padding a row carries. A realistic left-padded batch from the HF tokenizer, comparing each prompt alone against the same prompt inside the batch:

prompt	pads	HT drift	Bridge drift
"The capital of France is Paris and it is beautiful"	0	2.5e-06	0.00000
"Hello world"	8	4.3e-06	2.95003
"Machine learning models require careful evaluation"	4	1.8e-06	6.80700

This also breaks enable_compatibility_mode(), whose documented contract is "HookedTransformer-equivalent numerics". On unpadded input compat mode matches HT exactly (max |logit diff| 0.000e+00). Left-padded it does not: loss 4.503170 → 11.154946.

Root cause

transformer_bridge.py:1532 derives position_ids only for batched list input:

# Auto-compute attention_mask + position_ids for batched list input
# when the caller didn't supply them. Matches HF generation convention.
if (
    _is_batched_list
    and attention_mask is None
    ...
):
    ...
    position_ids = attention_mask.long().cumsum(-1) - 1
    position_ids.masked_fill_(attention_mask == 0, 1)
    kwargs["position_ids"] = position_ids

Pre-tokenized tensor input never reaches that branch, so HF falls back to a plain arange and the padding offset is never removed. Supplying position_ids by hand fixes it exactly — max |logit diff| drops from 1.361e+02 to 9.918e-05.

That also makes the bridge inconsistent with itself: the same batch gives different logits depending on whether you passed strings or token IDs. Max |logit diff| between the two input forms for identical content: 4.142e+01.

Code example

import torch
from transformer_lens import HookedTransformer
from transformer_lens.model_bridge import TransformerBridge

ht = HookedTransformer.from_pretrained("gpt2", device="cpu").eval()
br = TransformerBridge.boot_transformers("gpt2", device="cpu")
br.enable_compatibility_mode()
br.eval()

seq = ht.to_tokens("The capital of France is", prepend_bos=False)
L, n_pad = seq.shape[1], 3
pad = ht.tokenizer.eos_token_id

left = torch.cat([torch.full((1, n_pad), pad), seq], dim=1)
mask = torch.cat([torch.zeros(1, n_pad, dtype=torch.long),
                  torch.ones(1, L, dtype=torch.long)], dim=1)

with torch.no_grad():
    print("HT     unpadded :", ht(seq, return_type="loss").item())
    print("HT     left-pad :", ht(left, attention_mask=mask, return_type="loss").item())
    print("Bridge unpadded :", br(seq, return_type="loss").item())
    print("Bridge left-pad :", br(left, attention_mask=mask, return_type="loss").item())

# HT     unpadded : 4.503170
# HT     left-pad : 4.503169   <- invariant, correct
# Bridge unpadded : 4.503170
# Bridge left-pad : 11.154946  <- silently wrong

System Info

• Installed from source, at dev-4.x (f17dff30)
• OS: macOS 26.5 (arm64), CPU-only
• Python 3.12.13; torch 2.11.0; transformers 5.13.0

Reproduced on three independent paths — boot_native, boot_transformers, and enable_compatibility_mode() — and on two models (gpt2, distilgpt2). Deterministic across repeats (identical to 1e-9).

Additional context

Scope, and an honest caveat. Raw HuggingFace behaves the same way: GPT2LMHeadModel(left_tok, attention_mask=mask) without position_ids also drifts by 1.361e+02, and is correct once position_ids are passed. So the bridge is faithfully passing HF's convention through, and one could argue this is HF's contract rather than a bridge defect.

I think it is still a bug here, for three reasons:

1. HookedTransformer deliberately corrects it (pos_embed(tokens, pos_offset, attention_mask)), and the migration guide actively pushes users from HT to the bridge. Same code, different numbers, no warning.
2. The bridge already applies exactly this correction for list input, so the two input forms disagree with each other.
3. enable_compatibility_mode() documents "HookedTransformer-equivalent numerics", and this breaks that contract specifically.

Expected behaviour & fix pointer

When an attention_mask is supplied, position_ids are absent, and the mask indicates left padding, derive positions from the mask — the same two lines the list-input branch already uses:

if (
    attention_mask is not None
    and "position_ids" not in kwargs
    and not _is_inputs_embeds
    and attention_mask.ndim == 2
    and bool((attention_mask[:, 0] == 0).any())
):
    _pos = attention_mask.long().cumsum(-1) - 1
    kwargs["position_ids"] = _pos.masked_fill(attention_mask == 0, 1)

Locally this restores per-row parity with HT exactly (drift 0.00000 on every row of the batch above), with tests/unit/model_bridge + tests/unit/test_tokenizer_padding_side.py green (3955 passed, 27 skipped, 10 xfailed).

Interaction with #1607 / #1608 — these compound

Two independent bugs on the same path. Measured across four code states, gpt2, aggregate loss on the left-padded batch above (HT reference 4.814578):

state	batch loss	max |logit diff|	compat drift
dev-4.x today	7.254170	1.361e+02	6.651776
#1608 alone	7.411356	1.361e+02	11.115239
position fix alone	6.688155	9.918e-05	1.343308
both	4.814578	9.918e-05	0.000001

They separate cleanly: the position fix corrects the logits (1.361e+02 → 9.918e-05), #1608 corrects the loss aggregation, and only both together give correct batched loss.

Worth flagging that #1608 alone increases the reported drift (6.65 → 11.12) — two errors were partially cancelling. That is not a regression in #1608; it is an argument for landing both, and it would look like one if #1608 were benchmarked on padded batches in isolation.

Also: on the Native bridge, #1608's masked-softmax fix is a prerequisite for even observing this one — without it a fully-masked query row NaNs and the left-padded logits come back nan rather than merely wrong.

Deliberately not proposed here

generate() with left padding. HookedTransformer.generate has no attention_mask parameter (only **generation_kwargs), so a mask passed there is silently swallowed and I could not construct a valid comparison. There may or may not be a problem; I have no evidence either way and would rather not guess.

Also not proposed: changing the _is_batched_list branch itself. It is correct as written; this is about extending the same correction to tensor input.

Test coverage

No test asserts left-padding invariance for the bridge's forward pass, which is why this is invisible. A regression test is cheap and needs no real checkpoint — boot_native shows it — though a gpt2 test is the one that would have caught the compat-mode contract break.

Acceptance:

• position_ids derived from attention_mask for tensor input, matching the existing list-input behaviour
• Left-padding invariance regression test for forward / return_type="loss"
• Compat-mode parity test against HookedTransformer on a left-padded batch

Checklist

• I have checked that there is no similar issue in the repo (required) — searched left padding, position_ids, padding_side; nothing open or closed covers this