When developing our method, it felt weird and unprincipled to me to do a forward pass with two things but then only backprop to one. The gradient with respect to all parameters is the thing which we know succeeds at actually descending on an objective. Do we even know if learning happens properly when you mask out gradients to a bunch of parameters, and why should we think that this localizes the behavior to these parameters? As it turns out, all adapter fine-tuning (e.g., LoRA) makes a similar assumption
I don’t think this is quite right. LoRA is ordinary gradient descent on a single loss. GRAFT isn’t, and in principle it allows cycles and oscillations, which don't appear in ordinary gradient descent.
Gradient descent and LoRA. In gradient descent there is one loss L(θ), and every parameter moves downhill: θ' = −∂L/∂θ (in the small-learning-rate limit). This can’t cycle. Along the trajectory, L’ = −‖∂L/∂θ‖² ≤ 0, with equality only at stationary points, so the parameters can never return to where they started.
LoRA is a special case. Say the frozen parameters are θ, the trainable ones are u, and the loss is L(θ, u). Training u with θ fixed at θ* is just gradient descent on the single function u ↦ L(θ*, u). The frozen parameters are part of the loss. So LoRA can’t cycle or oscillate either.
GRAFT. Call the RETAIN parameters r and the FORGET parameters f. Let L_U be the average loss on unflagged episodes and L_D the average loss on flagged episodes. RETAIN learns only from unflagged episodes. FORGET learns from both, with flagged episodes upweighted by k:
r’ = −∂L_U(r, f)/∂r
f’ = −∂[L_U(r, f) + k·L_D(r, f)]/∂f
Following your work, we’ll use k = 2 throughout. In general, k should be chosen so that routing doesn’t change how strongly a flagged episode x is reinforced: (k − 1)·‖∂ log π(x)/∂f‖² = ‖∂ log π(x)/∂r‖².
Each block does honest gradient descent on its own loss, but the two losses differ. So the joint dynamics form a two-player game, not descent on one objective, and the argument that rules out cycles no longer applies.
This can go round in a circle. Take L_U(r, f) = ⟨r, f⟩ and L_D(r, f) = −⟨r, f⟩, where r and f are vectors of the same dimension. Then
r’ = −f, f’ = −∂(⟨r, f⟩ − 2⟨r, f⟩)/∂f = r.
The solution is r(t) = R cos t, f(t) = R sin t: a circle traversed forever, with ‖r‖² + ‖f‖² conserved.
Intuitively, RETAIN wants ⟨r, f⟩ small, so it wants r to point opposite to f: it runs away from FORGET. FORGET’s effective loss is −⟨r, f⟩, so it wants f to point along r: it chases RETAIN. One player evades and the other pursues, so they go round and round, like the game matching pennies. The cycling comes entirely from the two losses disagreeing about how the two blocks should relate to each other. Here, the unflagged data wants the adapters’ contributions to point in opposite directions, and the flagged data wants them aligned.
Discrete steps make it worse. With learning rate ε, one step maps (r, f) to (r − εf, f + εr). That multiplies the radius by √(1 + ε²) > 1 on every step, so in practice you get an outward spiral and unbounded parameter growth.
This can also occur with bounded loss. This example above uses unbounded losses, but bounded ones work too. With scalar r and f, L_U = tanh(r)·tanh(f) and L_D = −L_U, the losses stay in [−1, 1], and sinh²(r) + sinh²(f) is conserved, so the trajectories are again closed loops.
I see, interesting.
I suppose the thing I said in the post holds for just a single update of FORGET against a forward pass with both adapters active.
We did have an intuition that the forget and retain adapter could be caught in a "tug of war" because of the conflicting objectives. My intuition is that anchor training ought to prevent this in practice.
I've personally never made the intuition explicit like this, though; nice, clean example!
Your intuition about anchor training is correct. The anchor loss acts like "friction". Anchor episodes are sampled from the deployed model and only update RETAIN, so the anchor loss L_A depends on r alone:
r' = −∂L_U(r, f)/∂r − β·∂L_A(r)/∂r
f' = −∂[L_U(r, f) + 2·L_D(r, f)]/∂f
So L_A adds a damping term to RETAIN's own dynamics, wherever L_A is convex.
In the toy, take L_A = ½‖r‖². Then r' = −f − βr and f' = r.
• β = 0: the circle.
• 0 < β < 2: a damped spiral. Training still oscillates on the way, but it converges.
• β ≥ 2: it converges without oscillating.
With discrete steps, the squared radius is multiplied by 1 − εβ + ε² per step. So the anchor beats the outward spiral exactly when β > ε. That's a fairly mild requirement.
(NB: I had an intuition that GRAFT could lead to cycles because it had two components with different loss functions, like GAN. But Opus 5.5 came up with the examples.)
Thanks, this is great! Just wanted to note that the fact that the LoReFT approach appears to work so well is some evidence that the FORGET adapter probably functions similarly to a steering vector.
And also the approach makes some meaningful progress on this front:
> figuring out how to ablate the FORGET adapter without harming performance on target tasks would be particularly practically valuable.
Do a mech interp on the learned adapter
This seems especially useful. The FORGET adapter is "concentrated essence of reward hacker", which seems like a useful model organism and a fruitful target for mech interp: what layers and directions does it interact with? What in the main model is it turning up or down?
Localizing behaviors and controlling their deployment solves many of the same problems as unlearning
Somewhat, but I suspect it's a lot easier than unlearning. Most work in unlearning has been aimed at things originally learned by SGD. SGD produces a firehose of bits of supervision, but these are scattered across every topic/skill under the sun.
RL post-training produces very few bits of supervision, very narrowly focussed on what most improves improves performance against a specific task metric (some of which will generalize to varying extents, most parts will be task-specific). It tends to work only on things were there was already some >0 performance to uplift.
Gradient routing task-specific RL updates thus seems potentially vastly easier then doing do for specific tasks or data out of everything learned during SGD. Notoriously, unlearning that is hard: it's easy to damage capabilities but this almost invariably leaves a lot of the underlying circuitry/data intact, so relearning the skill/information becomes easy: the SGD-learned data we're trying to unlearn is easy to damage but hard to eradicate all traces of.
A proviso to that is that getting performance back down to the base model level is likely easy and robust, but getting it below that you are basically returning to thenfull unlearning SGD problem, which is known the be fragile/hard.
Ah, I was sort of unclear with this passage. I think GR in general (not our method) is a good method for unlearning (or never-learning) in pre-training. The idea is that this is better than post-hoc unlearning because you encourage circuitry to be built in some forget adapter. Then when you ablate the adapter, the circuitry is straight up not present in the deployed model (as opposed to e.g. gradient ascent or alignment SFT, where we suspect the circuit is merely suppressed, as you say).
I'm speaking as myself here and not on behalf of any of my co-authors. A bunch of the stuff in this post is explicitly speculative.
I've written this post as a way of centralizing some points I want to make in discussion/in response to others who've asked questions about the GRAFT paper. This is organized more like a series of notes; there's not a specific thesis in here, though I do make opinionated claims.
Specifically, a few people have asked what follow-up work I'd like to see; at the bottom of this post is a big brain-dump answer to that question.
Connections to other works
Selective modularity
Alex Cloud and Jacob G-W have already written an excellent treatise on this topic. I'll just emphasize a few select points.
GRAFT in the context of other gradient routing work
Some related ideas
Some looser associations
Practical Considerations
An aside on the objective
An optimistic speculation
Expand
It would be really great if various intuitions from persona selection, emergent misalignment, etc. held, and "misaligned behaviors" naturally clustered, especially as models scaled. Then this would be a more general solution to misalignment from RL; maybe this could rescue recent results that suggest that midtraining effects alone are inadequate.
However, the standard of evidence for success in this story should be extremely high: the effect should be really robust, and all of the concepts invoked here should be understood very systematically.
Soapbox: Alignment is still hard BTW
Expand
Follow-up work
My impression is that not a lot of work on gradient routing has been done overall, but it seems to have conceptual merit, so there are a lot of potential ideas in the space and there are many projects which would be high VoI on the margin.
Off the top of my head here are a bunch of possible questions to answer.
Scaling
Expand
It's a simple empirical question whether a given model behavior routes into the FORGET adapter in the way we'd intuitively like, and some interesting model behaviors only happen above a certain scale. On the flipside, there are a lot of design variables (see below) and it may be hard to study them systematically at scale, so it won't be easy to prove a negative.
Conceptual advances
Expand
Optimizing tradeoffs
Expand
Our setup still faces tradeoffs fundamentally associated with the gap between the training reward and the desired behavior at deployment. To illustrate them in our context: to the extent that we don't optimize our policy against our monitor, we are also exploring in a direction we know is suboptimal. In order to address this we need to train the deployed policy, but either we need to risk exposing it to the monitor again (that's what we do in our main small-scale setting), or we need to introduce external sources of good signal (that's what the anchor environment addresses). We can also attempt to explicitly train in an off-policy fashion, but that's cursed.
It's possible that in more realistic, frontier-like settings, some of these problems are less prevalent. For example, maybe only a small fraction of prompts in a small fraction of environments admit reward hacks, so the rest of the training serves as a natural anchor environment. I have roughly no intuition about what these factors balance out to, but it really feels like there should be some good engineery hacky ways to balance tradeoffs to our liking in realistic settings (more on this below).
Finally, note that the method will probably not be applied in practice if it comes at a significant capability cost, so figuring out how to ablate the FORGET adapter without harming performance on target tasks would be particularly practically valuable.
Method variations
Expand
The above points already suggest several possibilities for different things one could try:
Engineering advances
Expand
Setting variations
Expand
Some of our experimental design settings were a bit arbitrary, and could be studied systematically.
A mindset note
Expand
The number of moving parts may feel harrowing. If there's a kernel of insight in here, how are we to find it? At the beginning of our project, we had to make a lot of arbitrary design and parameter choices just to have a setting constrained enough to think about and gather evidence on. However, the flip side of this is that the method has a lot of affordances for potential development. In the bigger picture, I think this is a plus for gradient routing, and on net justifies more gradient routing work. Unfortunately, it may just be the case that most techniques that sound intuitively plausible don't work, and simultaneously, an actually good idea can take many tries before it works.
Anyone who wants to really dive in and could benefit from more context for their LLM agents can refer to this dev version of our code for the RL settings, with all the messiness and potential confusion that implies. In particular, results from some underpowered/exploratory experiments are found here.
AIUI, what they call "absorption" consists of both what we call "absorption" and what we call "inoculation", though they explicitly discuss the distinction between the underlying concepts. ↩︎
-
-
-
↩︎The total cost in FLOPs is the same or less on a per-sample basis.
We found in our experiments that gradient routing left sample efficiency unchanged on average, with a small minority of outliers in both directions.
A more specific concern is that throughput suffers because one can't batch routed backward passes with non-routed ones, but at large scales gradient accumulation becomes necessary anyway, so this factor becomes irrelevant. This issue can also be mitigated with a variety of tricks at smaller scale.