Deep dive into the ChoiceHead: task tokens, self-attention over choice embeddings, and bilinear scoring that is invariant to the number and order of choices.
After the Qwen backbone produces hidden states, the ChoiceHead receives three things extracted from that tensor:
answer_end_idx), which has attended to all choices.A task_type index per example also arrives, indicating whether this is a choice, noul, or score question. It is used to look up the task embedding.
The head prepends a task token to the choice embeddings before running self-attention. This token is a learned embedding indexed by task type (0=choice, 1=noul, 2=score). It tells the self-attention layers what kind of question is being answered, allowing the same head weights to handle all three task types.
task_token = self.task_embedding(task_type).unsqueeze(1) # [B, 1, H]
x = torch.cat([task_token, choice_embeddings, answer_embedding.unsqueeze(1)], dim=1)
# shape: [B, C+2, H] (task + C choices + answer)The padding mask passed to the self-attention layers marks the task token and answer token as always valid, while padding choice slots get a False (they contribute −∞ after masking and zero after softmax).
The head stacks num_layers identical AttentionBlock modules. Each block is a standard pre-norm transformer layer: LayerNorm → multi-head self-attention → residual, then LayerNorm → MLP → residual.
The attention inside the head projects from hidden dimension H down to new_dim (default 512) for Q, K, and V, then projects back to H via the output projection. This keeps the head compact relative to the backbone’s hidden dimension (1024 for Qwen3-0.6B).
Scores are computed as: QK² / sqrt(new_dim), masked with the padding mask, and softmaxed. Padded positions get −∞ before softmax so their attention weights are exactly zero.
After the self-attention stack, the head scores every choice using a bilinear interaction between the answer slot and each choice slot:
answer = self.answer_proj(x[:, -1]) # [B, D] — project last slot
choices = self.choice_proj(x[:, 1:-1]) # [B, C, D] — project choice slots
logits = einsum("bd,bcd->bc", answer, choices) / sqrt(new_dim)This is a dot product between the answer vector and each choice vector in a shared projection space of dimension D. The scaling by sqrt(new_dim) mirrors the scaling in standard attention.
Why bilinear? A dot product is symmetric and invariant to the number of choices — adding or removing a choice does not change the scores of the others. This makes the head applicable to any number of choices without retraining, and the final logits are masked for padding slots before cross-entropy.
When lora_last_k_layers > 0, PEFT wraps the last K Qwen layers with LoRA adapters targeting the four attention projection matrices (q_proj, k_proj, v_proj, o_proj). The frozen layers below are unmodified.
| Hyperparameter | Typical value | Effect |
|---|---|---|
| r (rank) | 8 | Number of singular vectors in the low-rank decomposition. Higher = more capacity. |
| alpha | 16 | Scaling factor: the effective LR for LoRA weights is lr * alpha/r. |
| dropout | 0.0 | Regularisation inside LoRA. Zero is standard for short runs. |
| last_k_layers | 12 of 20 kept | Layers 8–19 get adapters; layers 0–7 are fully frozen. |
The backbone requires_grad is False for all layers when lora_last_k_layers == 0 (head-only training). With LoRA, only the adapter weights are trainable in the backbone; the base weights stay frozen and never appear in optimizer state.
Say the answer out loud before revealing it — retrieving is what makes it stick.
What three things does the ChoiceHead receive from the backbone?
What is the task token and why is it prepended?
Why is bilinear scoring preferred over a linear classifier?
What happens to padding choice slots in the bilinear logits?