Problem
In mindspeed_mm/models/reward_model.py, the reward_token == "mean" path appears to treat sequence_lengths as a token count, but it is
actually computed as the last valid token index.
Relevant code:
sequence_lengths = torch.eq(input_ids, self.pad_token_id).int().argmax(-1) - 1
sequence_lengths = sequence_lengths % input_ids.shape[-1]
This makes sequence_lengths[i] the index of the final non-padding token.
Example:
input_ids: [10, 11, 12, PAD, PAD]
first PAD index: 3
sequence_lengths: 3 - 1 = 2
So sequence_lengths == 2 means the valid token range is indices 0..2.
However, the mean pooling branch uses:
valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1)
pooled_logits = torch.stack([logits[i, :valid_lengths[i]].mean(dim=0) for i in range(batch_size)])
Python slicing excludes the end index, so logits[i, :valid_lengths[i]] excludes the final valid token. For the example above, it
averages indices 0..1 instead of 0..2.
There is also a NaN risk when valid_lengths[i] == 0, because logits[i, :0] is empty and .mean(dim=0) returns NaN.
Expected Behavior
The mean pooling path should average all valid non-padding tokens, including the last valid token.
Problem
In mindspeed_mm/models/reward_model.py, the reward_token == "mean" path appears to treat sequence_lengths as a token count, but it is
actually computed as the last valid token index.
Relevant code:
sequence_lengths = torch.eq(input_ids, self.pad_token_id).int().argmax(-1) - 1
sequence_lengths = sequence_lengths % input_ids.shape[-1]
This makes sequence_lengths[i] the index of the final non-padding token.
Example:
input_ids: [10, 11, 12, PAD, PAD]
first PAD index: 3
sequence_lengths: 3 - 1 = 2
So sequence_lengths == 2 means the valid token range is indices 0..2.
However, the mean pooling branch uses:
valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1)
pooled_logits = torch.stack([logits[i, :valid_lengths[i]].mean(dim=0) for i in range(batch_size)])
Python slicing excludes the end index, so logits[i, :valid_lengths[i]] excludes the final valid token. For the example above, it
averages indices 0..1 instead of 0..2.
There is also a NaN risk when valid_lengths[i] == 0, because logits[i, :0] is empty and .mean(dim=0) returns NaN.
Expected Behavior
The mean pooling path should average all valid non-padding tokens, including the last valid token.