The problem: when reward models forget

It's well established that RLHF fine-tuning improves helpfulness scores on benchmarks while sometimes degrading other capabilities. But why does this happen at the representation level? Recent work on representation engineering[1] gives us tools to look inside, and what we find is geometric in nature.

Consider a language model's residual stream at layer \( \ell \). Before RLHF, the representations occupy a high-dimensional subspace. After training, they collapse toward a low-dimensional manifold aligned with the reward signal. The model hasn't "learned" to be helpful — it's learned to project everything onto the reward direction.

[1] Representation Engineering (Zou et al., 2023) — introduces methods for reading and controlling internal representations of neural networks.
$$ \text{rank}\left( \mathbf{H}^{(\ell)}_{\text{post}} \right) \ll \text{rank}\left( \mathbf{H}^{(\ell)}_{\text{pre}} \right) $$

where \( \mathbf{H}^{(\ell)} \in \mathbb{R}^{n \times d} \) is the matrix of hidden states at layer \( \ell \) across \( n \) inputs. The effective rank drops by 40–60% in the middle layers — precisely where we'd expect semantic composition to happen[2].

[2] Effective rank here is computed via the entropy of the singular value distribution, following Roy & Bhattacharya (2007).

Measuring collapse: the spectral lens

We can quantify this collapse by looking at the singular value spectrum of the representation matrix. A healthy representation has a gradual decay; a collapsed one has a sharp elbow. The following Python snippet computes the effective rank:

python effective_rank.py 
import numpy as np

def effective_rank(H):
    """Compute effective rank via singular value entropy."""
    s = np.linalg.svd(H, compute_uv=False)
    p = s / s.sum()
    p = p[p > 1e-10]  # filter near-zero
    entropy = -(p * np.log(p)).sum()
    return np.exp(entropy)

# Example: random vs collapsed
H_healthy = np.random.randn(100, 768)
H_collapsed = np.random.randn(100, 768) @ np.diag(
    np.exp(-np.linspace(0, 8, 768))
)
output

Visualizing the spectrum

The difference becomes immediately obvious when we plot the singular value distributions side by side. A pre-RLHF model retains representational capacity across dimensions; post-RLHF, the spectrum drops off sharply after the first few components[3].

[3] We normalize singular values to sum to 1 for comparison across models of different scales.
Fig 1. Normalized singular value spectrum, layer 16. Blue = pre-RLHF, purple = post-RLHF.
Fig 2. Cumulative variance explained. Post-RLHF reaches 95% with ~30 components vs ~120 pre-RLHF.
Fig 3. Effective rank by layer. Collapse is most severe in layers 12–20 (semantic composition zone).

Why does this matter for alignment?

Representation collapse isn't just a theoretical curiosity. It has direct consequences for alignment:

First, a model with collapsed representations loses its ability to make fine-grained distinctions. If safety-relevant features (deception detection, goal stability, situational awareness) live in the dimensions that get compressed, RLHF may silently destroy exactly the capabilities we need for scalable oversight.

Second, collapsed representations make probing less reliable[4]. If we train linear probes to detect deceptive behavior in a post-RLHF model, we may be fitting to the low-dimensional shadow of the original representation rather than the full feature. The probe appears to work on train data but fails on distribution shift — not because deception is hard to detect, but because the geometric structure we need is gone.

[4] This connects to the "probing illusion" — high probe accuracy doesn't imply the representation faithfully encodes the feature.

A formal characterization

Let \( \mathcal{M}_\theta \) be the manifold of internal representations parameterized by \( \theta \). RLHF training with reward model \( R \) induces a gradient flow that, in the limit, satisfies:

$$ \frac{d\theta}{dt} = \nabla_\theta \, \mathbb{E}_{x \sim \mathcal{D}} \left[ R\bigl(f_\theta(x)\bigr) \right] - \beta \, \nabla_\theta \, D_{\mathrm{KL}}\!\left(\pi_\theta \| \pi_{\mathrm{ref}}\right) $$

The KL penalty \( \beta \) controls how far the model drifts from the reference. When \( \beta \) is too small, the reward term dominates and representations collapse onto the reward manifold. When \( \beta \) is large, the model barely changes. The geometry of the transition between these regimes turns out to be a phase transition[5] — there is a critical \( \beta^* \) below which collapse is sudden and irreversible.

[5] We use "phase transition" loosely here — the effective rank as a function of β shows a sharp sigmoid, not a true discontinuity.

What can we do about it?

Three directions seem promising. First, monitoring effective rank during training as an early warning signal. If we see the rank dropping below a threshold in the middle layers, we can adjust \( \beta \) or pause training. Second, adding an explicit rank regularizer — penalizing low effective rank in the loss function. Third, and most speculatively, using representation engineering to "re-inflate" collapsed representations post-hoc, by identifying the lost directions and adding them back as controlled perturbations.

None of these are solved problems. But the geometric lens gives us something that reward curves alone don't: a way to see when the model is losing internal structure, even if benchmark scores keep going up.