[Bug Report] TransformerBridge seq2seq loss ignores labels and derives decoder inputs/targets from encoder input_ids

[emerardd]

Describe the bug

TransformerBridge.forward(..., return_type="loss") applies the shifted decoder-only causal loss to encoder-decoder models such as BART and T5.

For seq2seq models, the encoder input_ids and decoder labels are different sequences and may have different lengths. However, the current Bridge path:

1. derives decoder_input_ids from the encoder input_ids, even when labels are provided;
2. discards the loss already computed by the Hugging Face model;
3. recomputes a shifted causal loss against the encoder input_ids.

As a result:

• changing labels may not change either the returned Bridge loss or logits;
• the returned loss does not match Hugging Face seq2seq loss;
• source and target sequences of different lengths can raise a dimension error;
• encoder masking semantics may be incorrectly reused for decoder targets.

This is particularly risky because the returned scalar is finite and plausible when the source and target happen to have the same length, so the incorrect result can be used silently in training or interpretability experiments.

Code example

The following reproduction creates a tiny random BART model from config. It does not download any model or tokenizer from the Hub.

import torch
from transformers import BartConfig, BartForConditionalGeneration

from transformer_lens.model_bridge.sources._bridge_builder import (
    build_bridge_from_module,
)


torch.manual_seed(1234)

config = BartConfig(
    vocab_size=32,
    d_model=16,
    encoder_layers=1,
    decoder_layers=1,
    encoder_attention_heads=2,
    decoder_attention_heads=2,
    encoder_ffn_dim=32,
    decoder_ffn_dim=32,
    max_position_embeddings=32,
    pad_token_id=0,
    bos_token_id=1,
    eos_token_id=2,
    decoder_start_token_id=2,
)

model = BartForConditionalGeneration(config).eval()
bridge = build_bridge_from_module(
    model,
    "BartForConditionalGeneration",
    hf_config=config,
    tokenizer=None,
    model_name="tiny-random-bart",
)

source = torch.tensor([[4, 5, 6, 7, 2]])
labels_a = torch.tensor([[8, 9, 10, 2, -100]])
labels_b = torch.tensor([[11, 12, 13, 2, -100]])

with torch.no_grad():
    bridge_loss_a = bridge(source, labels=labels_a, return_type="loss")
    bridge_loss_b = bridge(source, labels=labels_b, return_type="loss")

    hf_loss_a = model(input_ids=source, labels=labels_a).loss
    hf_loss_b = model(input_ids=source, labels=labels_b).loss

    bridge_logits_a = bridge(source, labels=labels_a, return_type="logits")
    bridge_logits_b = bridge(source, labels=labels_b, return_type="logits")

    hf_logits_a = model(input_ids=source, labels=labels_a).logits
    hf_logits_b = model(input_ids=source, labels=labels_b).logits

print(f"Bridge loss A: {bridge_loss_a.item():.9f}")
print(f"Bridge loss B: {bridge_loss_b.item():.9f}")
print(
    "Bridge loss delta:",
    f"{torch.abs(bridge_loss_a - bridge_loss_b).item():.9f}",
)

print(f"HF loss A: {hf_loss_a.item():.9f}")
print(f"HF loss B: {hf_loss_b.item():.9f}")
print(
    "HF loss delta:",
    f"{torch.abs(hf_loss_a - hf_loss_b).item():.9f}",
)

print(
    "Bridge logits delta:",
    f"{torch.max(torch.abs(bridge_logits_a - bridge_logits_b)).item():.9f}",
)
print(
    "HF logits delta:",
    f"{torch.max(torch.abs(hf_logits_a - hf_logits_b)).item():.9f}",
)

short_decoder = torch.tensor([[2, 14, 15]])
short_labels = torch.tensor([[14, 15, 2]])

with torch.no_grad():
    hf_short_loss = model(
        input_ids=source,
        decoder_input_ids=short_decoder,
        labels=short_labels,
    ).loss

print(f"HF short-target loss: {hf_short_loss.item():.9f}")

with torch.no_grad():
    bridge(
        source,
        decoder_input_ids=short_decoder,
        labels=short_labels,
        return_type="loss",
    )

Observed output:

Bridge loss A: 3.489418030
Bridge loss B: 3.489418030
Bridge loss delta: 0.000000000

HF loss A: 3.463288307
HF loss B: 3.508223534
HF loss delta: 0.044935226

Bridge logits delta: 0.000000000
HF logits delta: 0.268047512

HF short-target loss: 3.496892214

RuntimeError: Size does not match at dimension 1 expected index [1, 4, 1]
to be no larger than self [1, 2, 32] apart from dimension 2

The exact random loss values may vary across dependency versions, but the important properties are:

• bridge_loss_a == bridge_loss_b;
• bridge_logits_a == bridge_logits_b;
• the corresponding Hugging Face loss and logits change with the labels;
• Hugging Face supports a decoder target length different from the encoder source length, while Bridge raises a shape error.

Expected behavior

For an encoder-decoder model, if labels are supplied:

• when decoder_input_ids are absent, the model should derive them by shifting labels to the right, following the Hugging Face seq2seq contract;
• return_type="loss" should compute loss against labels;
• the seq2seq loss should not apply an additional decoder-only logits[:, :-1] / labels[:, 1:] shift;
• positions where labels == -100 should be ignored;
• encoder and decoder sequences should be allowed to have different lengths;
• the encoder attention_mask should not be treated as the decoder target loss mask.

If return_type="loss" is requested for a seq2seq model without labels, the Bridge should either require labels explicitly or document another unambiguous target contract. It should not silently use encoder input_ids as decoder targets.

Root cause

The shared encoder-decoder path currently generates decoder_input_ids from encoder input_ids whenever the caller does not provide them:

https://github.com/TransformerLensOrg/TransformerLens/blob/f17dff30/transformer_lens/model_bridge/transformer_bridge.py#L1563-L1582

This condition does not exclude calls that provide labels. Supplying decoder_input_ids prevents Hugging Face seq2seq models from deriving decoder inputs from the labels.

After the driver forward pass, the Bridge extracts result.logits but does not use result.raw_output.loss. It passes the encoder input_ids to the common return finalizer:

https://github.com/TransformerLensOrg/TransformerLens/blob/f17dff30/transformer_lens/model_bridge/transformer_bridge.py#L1638-L1665

The finalizer then invokes the common shifted causal loss against those encoder tokens:

https://github.com/TransformerLensOrg/TransformerLens/blob/f17dff30/transformer_lens/model_bridge/bridge_core.py#L338-L347

https://github.com/TransformerLensOrg/TransformerLens/blob/f17dff30/transformer_lens/model_bridge/bridge_core.py#L350-L410

That objective is appropriate for decoder-only next-token prediction, but not for encoder-decoder teacher-forcing loss.

Affected architectures

The runtime reproduction above has been verified with BART.

The issue is in the shared TransformerBridge.forward() and loss finalization path rather than the BART adapter itself. Static inspection shows that the registered seq2seq adapters currently inherit supports_causal_loss = True, including:

• BART, MBART, M2M100, Marian, Pegasus, Blenderbot and LED;
• T5, MT5, legacy T5, LongT5 and SwitchTransformers;
• T5Gemma and T5Gemma2.

The centralized seq2seq architecture list is here:

https://github.com/TransformerLensOrg/TransformerLens/blob/f17dff30/transformer_lens/utilities/architectures.py#L9-L26

Only BART has been exercised in the reproduction above; the remaining families are listed based on the shared code path and adapter capability flags, not individual runtime verification.

Suggested direction

The preferred fix would be to add an explicit seq2seq loss path rather than reuse the causal loss path.

Possible behavior:

1. Detect encoder-decoder models before causal-loss finalization.
2. If labels are supplied and decoder_input_ids are not, allow the HF model to derive decoder inputs from the labels.
3. For scalar loss, return or otherwise preserve the architecture-correct loss from the raw model output.
4. For loss_per_token=True, compute unreduced cross-entropy directly against labels, without applying another causal shift, and ignore -100.
5. Require labels for seq2seq return_type="loss" / "both" unless another explicit target API is defined.
6. Keep encoder attention masks, decoder attention masks and loss-label masks as separate concepts.

A temporary safety fix would be to mark seq2seq adapters as not supporting the current shifted causal loss and raise an explicit error. That would be safer than returning a plausible but incorrect scalar, although full seq2seq support would be more useful.

It may also be clearer long-term to distinguish loss types explicitly, for example:

causal / seq2seq / unsupported

rather than representing all objective support through supports_causal_loss: bool.

Existing test/verification gap

I could not find a current test that compares:

bridge(
    encoder_input_ids,
    labels=target_ids,
    return_type="loss",
)

against the corresponding raw Hugging Face seq2seq loss.

The benchmark reference named hf_loss also appears to be populated from a Bridge call:

https://github.com/TransformerLensOrg/TransformerLens/blob/f17dff30/transformer_lens/benchmarks/main_benchmark.py#L1038-L1043

Therefore, existing Bridge loss-equivalence results do not appear to exercise the Hugging Face labels contract demonstrated above.

System Info

• Installation: source checkout, dependencies managed with uv
• OS: Windows, PowerShell
• Python: 3.12.10
• PyTorch: 2.11.0+cpu
• Transformers: 5.13.0
• Model used for reproduction: randomly initialized tiny BartForConditionalGeneration; no Hub download

Additional context

The same relevant path is present on dev-4.x at commit f17dff30.

Checklist

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

[emerardd]

@jlarson4, I'd like to work on this.

My proposed direction is to add an explicit seq2seq loss path: preserve Hugging Face's label-driven decoder input construction, use labels as targets without an additional causal shift, cover return_type="loss", "both", and loss_per_token=True, and add offline tiny BART and T5 parity tests, including different source/target lengths and -100 labels.

For seq2seq loss without labels, I plan to raise a clear error rather than silently fall back to encoder input_ids, unless maintainers prefer another contract.

Does that direction match the intended Bridge loss API? If so, could you assign this issue to me?

[jlarson4]

@emerardd I have assigned you, thank you for the thorough writeup. I am in agreement with your suggested scope, you are now assigned to work on this!

One related note: Decoder-only models also ignore labels. Tiny GPT-2 from config: bridge(inp, labels=nines, return_type="loss") == bridge(inp, return_type="loss") == 4.208895683, while HF returns 4.191130638 against the labels. Not as noticeable as seq2seq, but it is the same issue and a seq2seq fix should settle the decoder-only labels contract as well.