Senior Research Scientist at UK AISI working on frontier AI safety cases
Tomek Korbak
Karma: 662
Cool work! Reminds me a bit of my submission to the inverse scaling prize: https://tomekkorbak.com/2023/03/21/repetition-supression/
In practice I think using a trained reward model (as in RLHF), not fixed labels, is the way forward. Then the cost of acquiring the reward model is the same as in RLHF, the difference is primarily that PHF typically needs much more calls to the reward model than RLHF.
Thanks, I found the post quite stimulating. Some questions and thoughts:
-
Is LLM dynamics ergodic? I.e. is the time average equal to , the average page vector?.
-
One potential issue with this formalisation is that you always assume a prompt of size (so you need to introduce artificial “null tokens” if the prompt is shorter) and you don’t give special treatment to the token
<|endoftext|>. For me, it would be more intuitive to consider LLM dynamics in terms of finite, variable length, token-level Markov chains (until<|endoftext|>). While a fixed block size is actually being used during training, the LLM is incentivised to disregard anything before<|endoftext|>. So these two prompts should induce the same distribution:Document about cats.<|endoftext|>My name is;Document about dogs.<|endoftext|>My name is. Your formalisation doesn’t account for this symmetry. -
Dennett is spelled with “tt”.
-
Note that a softmax-based LLM will always put non-zero probability on every token. So there are no strictly absorbing states. You’re careful enough to define absorbing states as “once you enter, you are unlikely to ever leave”, but then your toy Waluigi model is implausible. A Waluigi can always switch back to a Luigi.
-
I don’t remember where I saw that, but something as dumb as subtracting the embedding of
<|bad|>might even work sometimes.
That’s a good point. But if you’re using a distilled, inference-bandwith-optimised RM, annotating your training data might be a fraction of compute needed for pretraining.
Also, the cost of annotation is constant and can be amortized over many training runs. PHF shares an important advantage of offline RL over online RL approaches (such as RLHF): being able to reuse feedback annotations across experiments. If you already have a dataset, running a hyperparameter sweep on it is as cheap as standard pretraining and in contrast with RLHF you don’t need to recompute rewards.
For filtering it was 25% of best scores, so we effectively trained for 4 epochs.
(We had different threshold for filtering and conditional training, note that we filter at document level but condition at sentence level.)
Good question! We’re not sure. The fact that PHF scales well with dataset size might provide weak evidence that it would scale well with model size too.
I’m guessing that poison-pilling the <|bad|> sentences would have a negative effect on the <|good|> capabilities as well?
That would be my guess too.
Have you tested the AI’s outputs when run in <|bad|> mode instead of <|good|> mode?
We did, LMs tends to generate toxic text when conditioned on
<|bad|>. Though we tended to have a risk-aversive thresholds, i.e. we used<|good|>for only about 5% safest sentences and<|bad|>for the remaining 95%. So<|bad|>is not bad all the time.Here it would be helpful to know what the AI produces when prompted by <|bad|>.
That’s a good point. We haven’t systematically investigate difference in capabilities between
<|good|>and<|bad|>modes, I’d love to see that.Just before public release, one could delete the <|bad|> token from the tokenizer and the model parameters, so switching to evil mode would require rediscovering that token embedding.
Yeah, you could even block the entire direction in activation space corresponding to the embedding of the
<|bad|>token
fixed, thanks!
I really liked the post and the agenda of improving safety through generative modelling is close to my heart.
we begin an online phase of its training: the agent starts acting in its environment and generating new task completions, which are recorded and fed back into the decision transformer as new training data
But you still need online access to our MDP (i.e. reward function and transition function), don’t you? And it’s access to MDP that drives novelty and improvement If you were just sampling whole trajectories from the model (asking the model itself to simulate reward function and transition model) and feeding them back into the model, you should expect any change (on average). Your gradients updates will cancel out, that’s a consequence of the expected-grad-log-prob lemma ().
It gets more nuanced when you account for doing ancestral sampling, but it adds problems, not solves them:
https://arxiv.org/abs/2110.10819Reproduce the “Learning to Summarize with Human Feedback” paper but with a frozen reward model.
On the other hand, in their follow-up work on instruction following, OpenAI claimed they used little online data (from fine-tuned policies):
https://arxiv.org/abs/2203.02155It would be interesting to figure out a way to factorize the policy in (a) over timesteps, i.e. produce distributions \) over actions conditional on partial trajectories
Levine derives that in his control-as-inference tutorial paper (section 2.3). Your expected exponential total reward is pretty close. Not that it acts a bit like an (exponentiated) Q function for your policy: it gives you exp-reward expected after taking action at state and following thereafter. The exponential works like a soft argmax, so it gives you something like soft Q-learning but not quite: argmax is also over environment dynamics, not only over policy. So it causes an optimism bias: your agent effectively assumes an optimal next state will sampled for it every time, however unlikely would that be. The rest of Levine’s paper deals with that.
good catch, yes, thanks!
Thanks for sharing your thoughts, I found these remarks extremely insightful!
It seems like ideal way forward is to more accurately capture what you actually care about, then optimize that—staying close to the original distribution feels like more of a hack to me. It seems like you view the original distribution of webtext as more principled or fundamental than I do, but I’m not sure what accounts for that difference.
A reply that comes to mins is that maybe being grounded in human knowledge, reasoning rules and values represented in web text has inherent value? Maybe web text is already approximately aligned with human preferences and you only want tweak that distribution a bit to match true human preferences? Assume that’s the case. Then, we can decompose LM alignment into (i) learning web text distribution and (ii) learning how to warp web text distribution. It seems that (ii) is easier than just learning aligned behaviour from scratch: your reward model doesn’t have to work well on arbitrary text but only text from distributions similar to webtext.
Another way of phrasing that point: maybe the assumption that you can have a perfect reward model is unrealistic and we can offload some of the complexity of learning a reward model to a prior given by web text? Or more philosophically, if you’re a Bayesian, you shouldn’t trust your reward model blindly, you should still have some prior.
Do you think these insights would generalise to the case where the language model may be interacting with some system during this fine-tuning phase? For example, if it generates queries to an external search engine or API, or has dialogue with a human, then the optimal policy is no longer equivalent to just generating the correct output distribution, as it now also involves environment observations.
That’s a good point and helps to make a distinction between generative models and policies. In the interactive case, your policy pi(a|s) is conditional distribution. You can equivalently view it as a collection of unconditional distributions {pi_s(a)}, one for each s, and for each of these you are likely to also have distribution collapse (single best action for a given state). Arguably, that’s what you want in RL.
So I think it mostly comes down to a philosophical difference. Do you want your LM to be a decision-maker acting in a world or a model of a some probability distribution over texts? If you want a decision-maker and training on language is just a scaffolding to get you there, maybe indeed staying close to the original distribution only has instrumental value?
But what if what you want is just an oracle-type conversational AI: a knowledge base and a common-sense reasoner. Maybe in this case staying close to human knowledge and inference rules represented in language is of inherent value?
I’m glad you found our post insightful!
I’m not sure what is the best energy allocation between modelling and inference here. I think, however, that the modelling part is more neglected (the target distribution is rarely even considered as something that can be written down and analysed). Moreover, designing good target distributions can be quite alignment-specific whereas designing algorithms for inference in probabilistic graphical models is an extremely generic research problem so we can expect progress here anyway.
I expect that in the current regime (only optimizing the policy a small amount), any method that does a reasonable job of maximizing reward while controlling how much the policy changes can be made to work in practice
Yes, that seems plausible. Though as you said, most methods that only change the policy a bit (early stopping, clipping in PPO) do that via implicit KL penalties and still can be seen as updating a prior.
there would be an exploration-exploitation trade-off, which is something that the RL perspective may again offer insight into.
Definitely exploration-exploitation issues could make the distribution collapse more severe and traditional RL tricks could help with that. But I still believe distribution collapse does not reduce to insufficient exploration and good exploration alone won’t solve it. In this specific instance, failing to find the optimal policy is not the problem, the optimal policy itself is the problem.
Fair point, I’m using “compositional” in an informal sense different from the one in formal semantics, closer to what I called “trivial compositionally” in this paper. But I’d argue it’s not totally crazy to call such preference models compositional and that compositionally here still has some resemblance to Montague’s account of compositionally as homeomorphism: basically, you have
get_total_score(response) == sum([get_score(attribute) for attribute in decompose(response)])