Prediction Is Not Output: One Head Cannot Do Both, Two Heads Can

Abstract

A language model uses one output head for two jobs: predicting what other people write, and producing its own replies. Where the text states a fact correctly, the two jobs ask for the same output. Where the people in the text are wrong, they do not: a good prediction reproduces their mistakes, and a good reply does not. We prove that one head trained on both jobs cannot do both well. Its best possible output is a weighted mixture of the two targets, which is wrong for both, whether the replies are trained by supervised fine-tuning, by reinforcement learning or by KL-regularized RLHF. We test this with a small transformer in a synthetic world where every target is known. A single head trained first to predict and then to answer, as language models are, forgets how to predict. A single head trained on both jobs at once answers correctly, but then expects other speakers to answer correctly 99% of the time, although only half of them do; when answering is given less weight, the same head instead answers the way the other speakers do, mistakes included. A model with two heads on a shared trunk, one trained to predict and one trained to answer, keeps both jobs right in every setting we tried. The price is speed: facts that appear only in plain statements reach the separate answering head more slowly.

Keywords