Supervised Adaptation for Stronger Decisions

Why People Wanted It

A language model can already choose among labeled options, but its ordinary answer-letter scores may not be enough for reliable decisions. The question is how to improve those decisions without replacing the pretrained model or turning the task into a separate reward-learning problem. The recipe is deliberately direct: use supervised examples, adapt the model with low-rank updates, and train a small decision head that starts from the model’s existing answer-letter weights.

The Core Move

Keep the language-model weights fixed, train LoRA adapters and an answer-letter decision head together, and minimize cross-entropy against the correct answer or a target distribution:

$\mathcal{L}=-\sum_{k=1}^{K} y_k\log p_k.$

Here $y_k$ is either a one-hot correct answer or a soft target distribution, and $p_k$ is the model’s predicted probability for option $k$.

Why the Head Starts from Answer-Letter Weights

Now the version for readers who want more of the machinery: the head is not an arbitrary classifier; its initial scores preserve the model’s own mapping from final hidden states to answer-letter logits. Let $h$ be the final normalized hidden state and $W_0$ the rows of the language-model output matrix for letters A through P. A trainable correction is added to those scores:

$z = W_0 h + A\,\mathrm{std}(h)+b.$

The original letter scores provide the starting point; $A$ and $b$ are the learned correction, initialized to zero. The feature standardization uses means and standard deviations computed once on training data and held fixed. The resulting logits are restricted to the valid options and normalized with softmax. For Gemma, the recipe preserves the original final-logit soft-capping, using $30\tanh(W_0h/30)$ in place of $W_0h$.

Why LoRA Changes the Representation

Here is the part that gets technical: the trainable head alone only reshapes fixed representations, whereas LoRA lets supervised examples alter how the model builds those representations. The recipe freezes the original language-model weights and adds rank-32 LoRA adapters, with $\alpha=64$ and no dropout, to attention and MLP projections; it also includes linear-attention projections in the Qwen models. The adapter updates and decision head are optimized jointly. In the reported Gemma 4 31B ablation, training only a head on frozen features scored 78.9 Capability, while LoRA with the head fixed at its original answer-letter initialization scored 84.9; the full recipe’s three-seed mean was 84.4. These are reported OpenDecisions results, not a guarantee for every model or dataset.

Why Option Shuffling and Soft Targets Matter

Now for the supervision details that change what the loss teaches: shuffling option order during training discourages a model from treating a particular answer letter as the answer, while soft targets preserve uncertainty or disagreement among labels instead of forcing every example into one hard choice. With a soft target, cross-entropy weights the log probability of each option by its target mass. The training pool included hard labels and soft labels from annotators or teacher probabilities; replacing soft targets with hard argmax labels produced no measurable Capability change in the reported Gemma ablation, so soft targets are a supported option, not a universally necessary ingredient.

Minimal Causal Code

This deliberately plain loop makes the supervised update visible. model_with_lora represents the frozen pretrained model plus trainable adapters; decision_head starts from the answer-letter weights plus a zero correction. The batch’s targets may be one-hot or soft distributions.

import torch
import torch.nn.functional as F

for batch in loader:
    # Randomly permute choices and apply the same permutation to targets.
    prompt, targets = shuffle_options(batch.prompt, batch.targets)

    h = model_with_lora(prompt).last_hidden_state[:, -1]
    logits = decision_head(h, valid_options=batch.valid_options)
    loss = -(targets * F.log_softmax(logits, dim=-1)).sum(dim=-1).mean()

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Quiz

What is the central training move in this recipe?

Fine-tune every pretrained weight and discard its answer-letter output behavior.

Use reinforcement learning and a reward model instead of labeled examples.

Replace the language model with a randomly initialized classifier.

Train LoRA adapters and a decision head jointly while keeping the original model weights frozen.

$ check

What You Gain, and What You Lose

You retain the pretrained model’s answer-letter behavior as the head’s starting point, while LoRA gives training a way to change internal representations. The method uses ordinary supervised learning and supports either correct-answer labels or probability targets. The tradeoff is that performance depends on the training examples and the starting checkpoint; a head trained only on frozen representations did not deliver the same improvement as adapting the model in the reported comparison. Soft labels and option shuffling are useful design choices, but the reported hard-label ablation did not show a measurable loss. The recipe is not a substitute for a strong instruction-tuned starting model: the reported pretrained-only Gemma 4 31B variant performed substantially below the instruction-tuned model’s trained result.

Evidence boundary

The exact recipe details and ablation figures above are reported for the specified training setup and Gemma 4 31B comparisons; they should not be read as universal effects. Calibration and the evaluation score are separate topics, not training mechanisms covered here.