[Bug Report] generate() cannot be given an attention_mask, so pre-padded input silently produces wrong continuations

[sohv]

Describe the bug

Neither TransformerBridge.generate() nor HookedTransformer.generate() takes an attention_mask parameter. When I pass an already-padded token tensor, there is no way to tell either method which tokens are padding, so the pads are treated as real context and the continuation diverges from the same prompt unpadded.

What makes this worse than a plain gap is that the mask is not rejected either. TransformerBridge.generate has a var-keyword named **multimodal_kwargs, so an attention_mask= I pass is silently absorbed into it and produces a third, different continuation with no warning. I checked, and the same happens for any unsupported kwarg — totally_bogus_kwarg=123 is accepted just as quietly.

generate(list_of_strings) is unaffected. I confirmed that path tokenises with forced left padding and builds the mask itself, so the longer prompt in a batch produces exactly the continuation it does on its own.

I found this while working on #1610, which fixed the equivalent problem in forward() by deriving position_ids from attention_mask. generate() cannot benefit from that fix, because the mask never reaches the forward it drives. I am raising it separately at @jlarson4's suggestion in that PR's review.

Code example

import torch
from transformer_lens.model_bridge import TransformerBridge

bridge = TransformerBridge.boot_transformers("distilgpt2", device="cpu")
bridge.eval()

tokens = bridge.to_tokens("The capital of France is")   # [1, 6]
n_pad = 4
pad_id = bridge.tokenizer.eos_token_id
padded = torch.cat([torch.full((1, n_pad), pad_id, dtype=tokens.dtype), tokens], dim=1)
mask = torch.cat(
    [torch.zeros(1, n_pad, dtype=torch.long), torch.ones(1, tokens.shape[1], dtype=torch.long)],
    dim=1,
)

kw = dict(max_new_tokens=6, do_sample=False, verbose=False)
print(bridge.generate(tokens, **kw)[0, tokens.shape[1]:].tolist())
print(bridge.generate(padded, **kw)[0, n_pad + tokens.shape[1]:].tolist())
print(bridge.generate(padded, attention_mask=mask, **kw)[0, n_pad + tokens.shape[1]:].tolist())

[262, 3139, 286, 262, 4141, 2066]   # unpadded — "the capital of the French Republic"
[262, 3139, 286, 262, 1578, 1829]   # padded, no mask — diverges at token 5
[262, 3139, 286, 262,  995,  338]   # padded + attention_mask kwarg — silently swallowed

All three should agree on the real tokens. I consider the third case the most misleading of them, because the caller supplied the correct mask and received neither an error nor the correct result.

System Info

• I installed transformer_lens from source (uv sync), on branch dev-4.x
• macOS (Darwin arm64)
• Python 3.12, transformers 5.13.0

Additional context

I see two directions here and do not have a strong preference between them. The first is to accept attention_mask as a real parameter on generate() and thread it into the per-step forward, which would let #1610's derivation apply. The second is to reject unknown kwargs rather than absorb them, so that at minimum the silent-wrong-answer case becomes a loud error.

Either way, I think tightening **multimodal_kwargs is worth considering on its own, since it currently swallows any unsupported kwarg and not just this one.

Per AGENTS.md §2, whatever lands on the Bridge should mirror to HookedTransformer.generate. It has the same gap through a differently named catch-all, **generation_kwargs, and I measured the identical divergence on pre-padded input: [262, 3139, 286, 262, 4141, 2066] unpadded against [262, 3139, 286, 262, 1578, 1829] padded.

I am happy to pick this up if you would like to assign it to me.

Checklist

• I have checked that there is no similar issue in the repo (required)

[jlarson4]

Looks good to me @sohv! Just a couple notes on my end:

I'd like to try accepting attention_mask as a real parameter. We can do an investigation into tightening **multimodal_kwargs in addition to that, but I would just be cautious that it does not have any knock-on effects (trimming kwargs we don't want trimmed, for example).

I recommend not mirroring to HookedTransformer, that AGENTS file should really be updated. We are in the process of deprecating HookedTransformer, it is likely any work done there will be dropped before it ever sees a release. The AGENTS file should really say "any changes to HookedTransformer need to be mirrored to the Bridge", but again at this point we aren't making any more additions to HookedTransformer, so the point is a bit moot.