adaptive optimizers and steering vector distillation

i just read an interesting paper on subliminal learning (arxiv:2606.00995v1) – it argues that subliminal learning is a special case of steering vector distillation (a student model trained on the outputs of a steered teacher model learns to imitate that steering).

the authors demonstrate this by examining the empirical activation similarity, which is the cosine similarity between the steering vector and the fine-tuning-induced shift in the student model's residual stream.

the paper concludes with two conditions under which steering vector distillation occurs. the optimizer has to be adaptive (adam) and weight updates have to be low-rank (lora). the paper discovers that the gradient on teacher-generated data carries a small component along the steering vector. this component is amplified by adaptive optimizers.

in this writeup, i want to provide some additional insights into why this might be the case.

gradients

let's begin by examining the student model's cross-entropy loss:

\mathcal L(\theta) = -\sum_x p_{\text{teacher}}(x\mid v)\log p_\theta(x)

where $p_\text{teacher}(\cdot | v)$ is the teacher's output distribution under steering vector $v$ .

we can rewrite:

p_\text{teacher}(x | v) = p_\text{teacher}(x) + \delta_v(x)

where $\delta_v(x)$ as the shift in token probabilities induced by $v$ .

through substitution, we get:

\mathcal L(\theta) = -\sum_x p_{\text{teacher}}(x)\log p_\theta(x) - \sum_x \delta_v(x)\log p_\theta(x)

taking the gradient w.r.t to $\theta$ gives us:

\nabla_\theta \mathcal L(\theta) = \underbrace{-\sum_x p_{\text{teacher}}(x)\nabla_\theta \log p_\theta(x)}_{\mathbf{g}_{\text{task}}} \underbrace{-\sum_x \delta_v(x)\nabla_\theta \log p_\theta(x)}_{\mathbf{g}_{\text{steer}}}

here, $\mathbf{g}_{\text{task}}$ represents the gradient of the kl divergence between the student model and the unsteered teacher model. this term optimizes for general language modeling capability. $\mathbf{g}_{\text{steer}}$ seeks to increase the log-prob of tokens related to the steering vector. this is because $\delta_v(x) > 0$ for such tokens, so $\delta_v(x)$ upweights the loss/gradient, and vice versa.

notice that $\delta_v(x)$ is small by construction since $v$ produces a subtle distributional shift. so, we can expect $\|\mathbf{g}_{\text{task}}\| \gg \|\mathbf{g}_{\text{steer}}\|$ .

one thing to note, however, is that taking steps towards $\mathbf{g}_{\text{steer}}$ doesn't guarantee that the student model's residual stream will steer towards $v$ . however, the authors provide empirical support, confirming that this is the case. when they computed the loss function's gradient w.r.t. the residual stream on steered data, they found that it had a more positive projection along $v$ than the gradient w.r.t. the residual stream on unsteered data.

optimizers

now, we can use our gradient deconstruction to understand why certain optimizers fail to elicit steering vector distillation.

let's begin with sgd. the update step is:

\bm \theta_{t+1} = \bm \theta_t - \eta \hat{\mathbf{g}}_t

where $\hat{\mathbf{g}}_t = \mathbf{g}_{\text{task}} + \mathbf{g}_{\text{steer}} + \bm{\epsilon}$ is the mini-batch estimate, which includes noise.

because $\|\mathbf{g}_{\text{task}}\| \gg \|\mathbf{g}_{\text{steer}}\|$ , picking a $\eta$ that is stable for $\mathbf{g}_{\text{task}}$ will be way too small to meaningfully move parameters in the $\mathbf{g}_{\text{steer}}$ direction. thus, sgd struggles to support steering vector distillation.

on the other hand, adaptive methods like adam/rmsprop don't suffer from this issue. let's ignore first moment estimates (momentum) for now, and examine their usage of second moment estimates:

\mathbf{\hat{v}}_t = \frac{\beta_2 \mathbf{\hat{v}}_{t-1} + (1-\beta_2)\mathbf{\hat{g}}_t^2}{1 - \beta_2^t}

and updates:

\bm \theta_{t+1} = \bm \theta_t - \eta \cdot \frac{\mathbf{\hat{g}}_t}{\sqrt{\mathbf{\hat{v}}_t} + \epsilon}

by using $\hat{\mathbf{v}}_t$ , adam/rmsprop perform coordinate-wise normalization. coordinates $\bm{\theta}_i$ with historically small gradients receive larger effective steps, and vice versa. empirically, this seems to prevent $\mathbf{g}_{\text{steer}}$ from being drowned out by a much larger $\mathbf{g}_{\text{task}}$ . the key is that, after normalization, coordinates where the steering signal is relatively strong can still affect the update, even when the raw gradient is dominated by $\mathbf{g}_{\text{task}}$ .

the paper runs multiple experiments which confirm this belief. the most straightforward ones compare adam/rmsprop to sgd/sgd with momentum. a more interesting experiment involved preserving the bottom 10% of $\mathbf s_i = \frac{1}{\sqrt{\mathbf{\hat v_{t,i}}}+\epsilon}$ values, while replacing the other 90% with the geometric mean of their $\mathbf s_i$ values. this achieves steering vector distillation that's on par with adam, suggesting that adam’s main advantage is not precise coordinate-wise scaling everywhere, but rather suppressing the small subset of high-gradient coordinates that would otherwise dominate.

low-rank updates

the paper doesn't provide a mechanistic hypothesis for why low-rank updates are necessary for steering vector distillation to occur. here's my tentative one:

we can conceptualize lora as restricting gradient updates to a low-dimensional subspace:

\mathbf P \hat{\mathbf{g}}_t = \mathbf P \mathbf{g}_{\text{task}} + \mathbf P \mathbf{g}_{\text{steer}} + \mathbf P \bm{\epsilon}

assuming the steering update is low-rank, $\mathbf P \mathbf g_{\text{steer}} \approx \mathbf g_{\text{steer}}$ . in addition, we can decompose $\bm \epsilon$ into parallel and perpendicular components w.r.t. the subspace:

\|\bm \epsilon\|^2 = \|\bm \epsilon_{\parallel}\|^2 + \|\bm \epsilon_{\perp}\|^2 = \|\mathbf P \bm \epsilon\|^2 + \|\bm \epsilon_{\perp}\|^2

so $\|\mathbf P \bm \epsilon\| \ll \|\bm \epsilon\|$ .

let a layer's residual stream value be $\mathbf h = \alpha \mathbf v + \mathbf z$ , where $\|\mathbf v\| = 1$ and $\mathbf v \perp \mathbf z$ . it might seem like if $\bm \epsilon$ changes $\mathbf h$ in direction $\mathbf z$ , this wouldn't matter because it doesn't affect $\alpha$ . but the residual stream gets normalized through something like rmsnorm in subsequent layers, where $\tilde{\mathbf h} = \frac{\mathbf h}{\|\mathbf h\|}$ , and the effective strength becomes:

\mathbf{v}^\top \tilde{\mathbf{h}} = \mathbf{v}^\top \frac{\alpha \mathbf{v} + \mathbf{z}}{\sqrt{\alpha^2 + \|\mathbf{z}\|^2}} = \frac{\alpha}{\sqrt{\alpha^2 + \|\mathbf{z}\|^2}}

so if $\bm \epsilon$ pushes the residual stream towards $\mathbf z$ , it can reduce the effective strength of $\mathbf v$ , even though $\alpha$ itself doesn't change.

projection removes $\bm \epsilon_{\perp}$ , . the noise that remains is $\epsilon_{\parallel}$ , which lies in the same low-dimensional update subspace as $\mathbf g_\text{steer}$ . this noise can still perturb $\alpha$ , but it's less damaging than perturbing $\mathbf z$ , because changes in $\alpha$ affect both the numerator and denominator of $\mathbf{v}^\top \tilde{\mathbf{h}}$ , while changes in $\mathbf z$ affect just the denominator.

tldr; lora reduces the mini-batch noise's dimensionality, which stops perpendicular updates from diluting the signal from $\mathbf g_\text{steer}$ . adam alone doesn't solve this because it still updates in all coordinates, including $\bm \epsilon_{\perp}$ , which weaken the signal of $\mathbf g_\text{steer}$ .