Reward Modeling › Reward Model Architecture

The standard approach to building a reward model starts with a pretrained language model and adds a simple regression head that maps the final hidden states to a scalar reward. This design philosophy reflects a key insight: rather than learning to understand language from scratch, we can use the rich representations already encoded in models that have been trained on large text corpora. The pretrained model provides the linguistic foundation, while the new regression head learns to interpret those representations through the lens of human preferences.

Think of the architecture as dividing the problem into two parts. The pretrained backbone is responsible for language understanding: it reads the prompt and response, parses the argument structure, identifies factual claims, and builds rich contextual representations of every token. The value head is responsible for quality assessment: given those rich representations, it extracts the signal most relevant to human preferences and collapses it to a single number. The backbone does the hard language work; the value head does the preference judgment. This division of labor is what makes the approach so sample efficient. Training the value head from scratch would require enormous amounts of preference data to simultaneously learn language understanding and preference assessment. By reusing a pretrained backbone, we only need the preference data to fine-tune the judgment layer.

The key insight is that the backbone's representations already encode many of the features that correlate with response quality: coherence, factual consistency, appropriate register, completeness of explanation. The value head's task is not to discover these features from scratch but to learn which combinations of already-discovered features predict human preference in the training domain.

The pretrained backbone is typically frozen or lightly fine-tuned during reward model training. Freezing the backbone is more computationally efficient and preserves the rich general language understanding. Light fine-tuning (using a very small learning rate for the backbone and a larger one for the value head) can improve performance when the training domain is significantly different from the pretraining corpus. In practice, most implementations fine-tune the entire model with a conservative learning rate, letting the backbone to gradually adapt while preserving its core representations.

Reward models typically use the same architecture family as the policy model they will train. If you're using RLHF to fine-tune a 7B parameter LLaMA model, your reward model might be initialized from the same pretrained checkpoint or a similar model in the family. This architectural alignment ensures the reward model can process the same input representations and has similar language understanding capabilities. The choice is not arbitrary: when the reward model shares the same "vocabulary" of internal representations as the policy model, it can more accurately evaluate the subtle qualities of responses that the policy might generate.

There is a meaningful tradeoff around reward model size. A larger reward model generally has better language understanding and can detect more subtle quality distinctions. However, it is also more expensive to run during the RL training loop, where it must be queried for every sampled response. In practice, teams often use a reward model that is the same size as or smaller than the policy, accepting some quality loss for the sake of training throughput. Some systems use reward model distillation to compress a high-quality large reward model into a smaller one that can be evaluated more cheaply during RL.

The key modification is replacing or augmenting the language modeling head. Instead of predicting the next token, we need to output a single scalar value that represents the quality of the entire response. This transformation converts a generative model into an evaluative one, shifting from the question "what comes next?" to "how good is this?" Because the parameters encoding language understanding are already in place, this transformation requires relatively little new learning. The backbone simply needs to route its existing representations through a new output projection.

One important architectural consideration is whether to initialize the reward model from the same checkpoint used to start the supervised fine-tuning (SFT) stage, or from the SFT model itself. In most RLHF pipelines, the SFT stage produces a model that has already learned to follow instructions and produce helpful responses. Initializing the reward model from this SFT checkpoint means the backbone already understands the response format expected by the system, potentially which makes it easier to learn quality distinctions within that format.

The core architectural question for reward modeling is this: how do we collapse an entire sequence of hidden states, one for each token in the prompt and response, into a single number that captures the overall quality? The solution involves selecting a representative hidden state and projecting it down to a scalar through a learned transformation.

Think of the sequence of hidden states as a series of snapshots taken at each position as the model reads through the text. Each snapshot captures what the model "knows" at that point in the sequence, accumulated from all prior tokens via the attention mechanism. The final snapshot, taken after the model has processed the entire sequence, is the most information-rich because it has had access to everything that came before. Using this final snapshot as the input to the value head is like asking someone to score a piece of writing after they have read the whole thing, rather than after each sentence.

The reward model architecture can be expressed as:

$$r(x, y) = f_\theta(h_{\text{final}})$$

where:

• $r(x, y)$: the scalar reward score output by the model
• $x$: the prompt text input to the model
• $y$: the response text to be scored
• $h_{\text{final}}$: the hidden state vector at the final token position
• $f_\theta$: the learned function (value head) mapping the hidden state to a scalar

This formulation captures the needed transformation at the heart of reward modeling. The input is a high-dimensional hidden state vector, perhaps 768 or 4096 dimensions depending on the model size, and the output is a single real number. The function $f_\theta$ must learn to extract and weigh the relevant features from this rich representation to produce a meaningful quality score.

This projection is typically implemented as a linear layer:

$$f_\theta(h) = w^T h + b$$

where:

• $f_\theta(h)$: the scalar output of the value head
• $w \in \mathbb{R}^d$: the learned weight vector for the value head
• $w^T$: the transpose of $w$, letting the dot product with $h$
• $h$: the input hidden state vector
• $b \in \mathbb{R}$: the learned bias term
• $d$: the hidden dimension size of the base transformer model

The linear layer performs a weighted sum over all dimensions of the hidden state. Each component $w_i$ learns to assign importance to the corresponding dimension $h_i$ of the hidden representation. Positive weights mean that feature contributes positively to the reward, negative weights indicate a negative contribution, and weights near zero suggest the feature is irrelevant for quality assessment. The bias term $b$ shifts the overall reward scale, letting the model to center its predictions appropriately.

Some implementations use a two-layer MLP for the value head rather than a single linear layer:

$$f_\theta(h) = w_2^T \cdot \text{ReLU}(W_1 h + b_1) + b_2$$

where:

• $W_1 \in \mathbb{R}^{d_{\text{hidden}} \times d}$: the first projection matrix, mapping the transformer hidden state to an intermediate dimension $d_{\text{hidden}}$
• $b_1 \in \mathbb{R}^{d_{\text{hidden}}}$: the bias for the first linear layer
• $\text{ReLU}(\cdot)$: the rectified linear activation, introducing nonlinearity between the two layers
• $w_2 \in \mathbb{R}^{d_{\text{hidden}}}$: the final projection vector collapsing the intermediate representation to a scalar
• $b_2 \in \mathbb{R}$: the output bias

The nonlinear value head allows the model to learn more complex combinations of the backbone's features. In practice, the single linear layer often performs comparably to the MLP, especially when the backbone is large and already captures high-level quality features. The additional parameters of the MLP also increase the risk of overfitting on small preference datasets.

Why use the final token position? In autoregressive models, information flows from left to right through causal attention. Each token can only attend to tokens that came before it, creating a natural accumulation of information as we move through the sequence. The final token's hidden state has "seen" all preceding tokens in both the prompt and response, which makes it a natural summary of the entire sequence. By the time we reach the end, the model has processed every word, every argument, and every nuance. This is analogous to using the [CLS] token representation in BERT-style models, which we covered in Part XXIV, where a special token is positioned to aggregate information from the entire input.

Some implementations average over all response token positions instead:

$$h_{\text{avg}} = \frac{1}{|y|} \sum_{t \in \text{response}} h_t$$

where:

• $h_{\text{avg}}$: the mean pooled hidden representation
• $|y|$: the number of tokens in the response
• $h_t$: the hidden state at token position $t$
• $\sum_{t \in \text{response}}$: a summation over all token positions belonging to the response

This mean pooling approach treats all positions as equally important and computes their centroid in hidden space. The intuition is that every part of the response matters, and averaging captures the "typical" representation across the sequence. However, using the final token is more common in practice because it naturally captures the complete context and requires no additional computation. The causal attention mechanism has already done the work of aggregating information, so the final position provides a ready-made summary. A potential advantage of mean pooling is greater stability when the most important content appears in the middle rather than the end, but in practice this benefit is small relative to the conceptual simplicity of the final-token approach.

A third variant, used in some systems, extracts the hidden state at the position corresponding to a special end-of-response token explicitly appended to every input. This ensures there is always a dedicated "summary" token regardless of what appears last in the tokenized sequence, which can improve consistency when response lengths vary widely.

The reward model must handle prompt-response pairs of varying lengths. Some prompts are brief questions, others are lengthy instructions with context. Similarly, responses range from terse answers to elaborate explanations. The architecture must gracefully accommodate this variability while maintaining consistent scoring semantics.

Think of the variable-length input problem as managing a sliding measuring tape. The reward always needs to be extracted from the right end of the tape, but the tape itself is a different length every time. The solution is to track where the tape ends and always read the final marker, regardless of the overall length.

The input is formatted as a concatenation of the prompt and response:

The model processes this sequence through the transformer layers, and we extract the hidden state at the EOS (end-of-sequence) token position for the value head. This approach uses the causal attention mechanism to ensure the reward is computed based on the complete context. The EOS token is a natural boundary marker, signaling where the response ends and giving a consistent extraction point regardless of the sequence length.

In practice, variable-length inputs require careful handling of attention masks during batching. When sequences in a batch have different lengths, shorter sequences are padded to match the length of the longest sequence in the batch. Padding tokens should not contribute to the final representation, so we must identify the position of the actual last token (not the last padding token) for each sequence. This is typically implemented by summing the attention mask to find the sequence length, then using that length to index into the hidden states.

The maximum sequence length is another practical constraint. Most transformer architectures have a fixed context window (2048 tokens for earlier models, 4096 or more for recent ones). Prompt-response pairs that exceed this limit must be truncated. The standard approach truncates from the right, which means very long responses may be partially cut off. Some implementations preferentially truncate the prompt rather than the response, reasoning that the response quality is what we are scoring. Choosing where to truncate involves a tradeoff: truncating responses can cause the model to miss important quality signals, while truncating prompts can deprive the model of context needed to judge appropriateness.