1The problem: choices must not see each other

To score every choice fairly, each one must be conditioned on the question but not on any sibling choice. If choice B could attend to choice A, the model could exploit their relative order or wording rather than their individual merit. The solution is a carefully crafted sequence layout combined with a custom attention mask.

2Token layout

The full input sequence is assembled from four parts in order:

The code reads the hidden state at the last content token of each choice, not at the closing > of </option>. This is intentional: the last content token carries the richest contextual representation for that choice, already conditioned on the full choice text.

# choice_read_idx is choice_end_idx minus the length of the "</option>" token sequence
choice_end_token_len = len(tokenizer.encode("</option>"))
choice_read_idx = [end - choice_end_token_len for end in span["choices_end_idx"]]

3Parallel position IDs

Standard position IDs increment monotonically from 0. In JEV, the question tokens still get sequential positions 0, 1, 2, … but every choice block restarts from the same position as the first choice. The answer tokens then start right after the longest choice, so every choice token is already in the past from the answer’s perspective.

Why does this matter? Rotary positional embeddings (RoPE, used in Qwen) encode the relative distance between tokens. If both choices start at position 51, their internal distance structure is identical and neither has a positional advantage over the other. A shorter choice’s last token lands at a lower absolute position, but the answer token is placed after the longest choice, so all choices are equally in the answer’s past.

4The attention mask

The attention mask is an L×L matrix (L = total sequence length) where a 1 means the row token can attend to the column token. Four regions get special treatment:

In code, the mask is built as a single L×L dense boolean tensor, then passed into the backbone as the attention_mask argument. Qwen’s flash attention path interprets it correctly.

5Putting it all together

The output of encode_example is a dictionary of tensors that goes straight to the DataLoader and then to the network:

{
  "input_ids":      [Q... <opt>C1</opt> <opt>C2</opt> ans...],
  "position_ids":   [0 1 2 ... | 51 52 53 | 51 52 53 | 54 55],
  "attention_mask": <L x L bool tensor>,
  "choice_read_idx":[idx_of_last_content_token_per_choice],
  "answer_end_idx": idx_of_last_answer_token
}

The choice_mask (which choices are real vs padding, used to batch variable-length examples) is computed later during collation. No labels are stored here; the dataset returns them separately and they are moved to the device in the training loop.