Encoder repetition penalty uses incorrect source rows after beam or sampled-return expansion

System Info

• Transformers: 5.19.0.dev0, main commit d6c1e71bd717bf092f8293f0c3c9bd4a5ac5401a; also reproduced on stable 5.18.0.
• Python: 3.12.14
• PyTorch: 2.14.1+cpu
• Accelerate: 1.15.0
• Platform: Linux x86-64, CPU, single process.
• Tiny models constructed from configuration; no model downloads or custom logits processors.

Who can help?

@Cyrilvallez

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

With encoder_repetition_penalty != 1, generation retains the original source IDs inside the processor while expanding score rows for beams or returned samples. Some hypotheses receive another source's token preferences, and later hypotheses receive no encoder penalty at all.

The following script first shows the row mismatch directly, then checks the processed scores from the public generate() API against an independent per-source calculation.

import torch
import transformers
from transformers import GPT2Config, GPT2LMHeadModel
from transformers.generation.logits_process import EncoderRepetitionPenaltyLogitsProcessor

torch.set_num_threads(1)
torch.manual_seed(123)
print("Transformers:", transformers.__version__)

source = torch.tensor([[1, 3], [4, 5]])
scores = torch.arange(1, 9, dtype=torch.float32).repeat(4, 1)
actual = EncoderRepetitionPenaltyLogitsProcessor(2.0, source)(
    torch.zeros(4, 1, dtype=torch.long), scores
)
print("Direct scores:", actual.tolist())

model = GPT2LMHeadModel(GPT2Config(
    vocab_size=8, n_positions=16, n_embd=8, n_layer=1, n_head=2,
    bos_token_id=1, eos_token_id=None, pad_token_id=0,
)).eval()

for mode in ("greedy", "beam", "sample_multiple"):
    options = {}
    copies = 1
    if mode == "beam":
        options = {"num_beams": 2}
        copies = 2
    elif mode == "sample_multiple":
        options = {"do_sample": True, "num_return_sequences": 2, "top_k": 0}
        copies = 2

    output = model.generate(
        source,
        attention_mask=torch.ones_like(source),
        max_new_tokens=1,
        encoder_repetition_penalty=2.0,
        return_dict_in_generate=True,
        output_scores=True,
        output_logits=True,
        **options,
    )
    # Beam processors receive log-probabilities; other modes receive logits.
    raw = output.logits[0]
    if mode == "beam":
        raw = raw.log_softmax(-1)

    expected = raw.clone()
    expanded_source = source.repeat_interleave(copies, dim=0)
    for row, tokens in enumerate(expanded_source):
        for token in tokens.unique().tolist():
            value = raw[row, token]
            expected[row, token] = value / 2.0 if value < 0 else value * 2.0

    print(mode, "correct:", torch.allclose(output.scores[0], expected))
    print("Max error per row:", (output.scores[0] - expected).abs().amax(-1).tolist())

For the direct example, each score row represents one of two hypotheses per source. Expected:

source 0, hypothesis 0: [1, 4, 3, 8,  5,  6, 7, 8]
source 0, hypothesis 1: [1, 4, 3, 8,  5,  6, 7, 8]
source 1, hypothesis 0: [1, 2, 3, 4, 10, 12, 7, 8]
source 1, hypothesis 1: [1, 2, 3, 4, 10, 12, 7, 8]

Actual on main and stable 5.18.0:

row 0: [1, 4, 3, 8,  5,  6, 7, 8]  # correct
row 1: [1, 2, 3, 4, 10, 12, 7, 8]  # uses source 1 instead of source 0
row 2: [1, 2, 3, 4,  5,  6, 7, 8]  # unchanged
row 3: [1, 2, 3, 4,  5,  6, 7, 8]  # unchanged

greedy correct: True
beam correct: False
sample_multiple correct: False

This also reproduces with a tiny T5 encoder-decoder model. In a three-token T5 beam-generation check, the second source's output changes from [0, 0, 0, 0] to [0, 4, 4, 4] after the source mapping is corrected, using identical model weights.

Expected behavior

Every beam or returned sample should receive the encoder repetition transformation for its own source. Batching and hypothesis expansion should preserve source ownership and apply the configured penalty to every hypothesis.

The processor documentation describes rewarding or penalizing tokens present in the input prompt. Generation expands sources with repeat_interleave, so each repeated hypothesis should retain its source's token IDs.

Suspected cause and possible fix

The processor caches the original B-row encoder_input_ids, then uses that smaller index tensor in gather/scatter against B × expansion score rows. PyTorch accepts the smaller index shape, so only the first B rows are modified.

A local correction expands the source IDs inside __call__:

encoder_input_ids = self.encoder_input_ids.repeat_interleave(
    scores.shape[0] // self.encoder_input_ids.shape[0], dim=0
)

The existing gather and scatter then use this local encoder_input_ids. This adds three lines and updates two existing lines. It preserves the no-expansion behavior and does not mutate the cached original source IDs.

Validation and related work

• 12 focused cases fail before the correction and pass afterward; six no-expansion controls pass before and afterward. Coverage includes batches of 1/2, expansion factors of 1/2/3, penalty factors of 0.5/2.0, positive/negative scores, and GPT-2/T5 beam and multiple-return sampling paths.
• All 48 existing tests in tests/generation/test_logits_process.py pass with this correction and a separate EOS-decay correction applied.
• The existing test_encoder_repetition_penalty_dist_process covers equal numbers of source and score rows; it does not exercise hypothesis expansion.
• The original implementation in #18675 already uses the unexpanded source IDs. #48642 concerns normalized repetition-penalty semantics and retains this row-mapping problem; this reproduction checks the appropriate score space for each decoding mode.
• Searches of issues and PRs did not find an equivalent source-row mapping report or active fix.

I can submit a PR with a focused fix and regression tests if this approach looks appropriate.

AI assistance was used to investigate this behavior and prepare this report. The reproduction and local correction were executed and verified in the environments above.