[Bug Report] TransformerBridge causal loss ignores attention_mask

[emerardd]

Describe the bug

TransformerBridge.forward(..., return_type="loss") and return_type="both" ignore the supplied attention_mask when computing causal LM loss.

The mask is forwarded to the underlying model, but BridgeCore._finalize_return() calls self.loss_fn(logits, input_ids, ...) without passing it. Padding transitions therefore contribute to both the loss sum and denominator.

Changing only masked padding token IDs changes the reported loss, while the correct mask-aware loss remains identical.

Code example

import torch

from transformer_lens.config import TransformerBridgeConfig
from transformer_lens.model_bridge import TransformerBridge

cfg = TransformerBridgeConfig(
    d_model=32,
    d_head=8,
    n_heads=4,
    n_layers=2,
    n_ctx=6,
    d_vocab=32,
    d_mlp=64,
    act_fn="gelu",
    normalization_type="LN",
    seed=7,
    initializer_range=0.2,
)
bridge = TransformerBridge.boot_native(cfg)

mask = torch.tensor([
    [1, 1, 1, 0, 0, 0],
    [1, 1, 1, 1, 1, 1],
])
inputs = {
    "pad_a": torch.tensor([
        [1, 2, 3, 0, 0, 0],
        [4, 5, 6, 7, 8, 9],
    ]),
    "pad_b": torch.tensor([
        [1, 2, 3, 31, 30, 29],
        [4, 5, 6, 7, 8, 9],
    ]),
}

for name, tokens in inputs.items():
    with torch.no_grad():
        logits = bridge(tokens, attention_mask=mask, return_type="logits")
        actual = bridge(tokens, attention_mask=mask, return_type="loss")
        expected = bridge.loss_fn(logits, tokens, attention_mask=mask)
    print(name, "forward:", actual.item(), "masked:", expected.item())

Output:

pad_a forward: 3.550006866455078 masked: 3.8386409282684326
pad_b forward: 4.128363132476807 masked: 3.8386409282684326

Expected behavior

return_type="loss" and return_type="both" should pass the forward-pass attention_mask into loss_fn(). Masked padding token IDs must not affect the result.

The same mask should also be applied when loss_per_token=True.

System Info

• Installed from source
• Branch: dev-4.x
• Commit: f17dff30
• Windows
• Python 3.12.10
• PyTorch 2.11.0+cpu
• transformers 5.13.0

Additional context

BridgeCore.loss_fn() and lm_cross_entropy_loss() already accept an attention_mask; it is lost between TransformerBridge.forward() and _finalize_return().

Native Bridge also exposes a related left-padding edge case: a fully masked padding query softmaxs to NaN, and those NaNs can poison later layers. Multiplying masked losses by zero is insufficient because NaN * 0 remains NaN. A complete fix needs NaN-safe masked loss handling and finite attention patterns for fully masked Native queries.

Checklist

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

[emerardd]

I'd like to work on this. I have a local fix with regression tests and will open a draft PR shortly.