Forecasting an AI model’s answer activations before it starts writing

Predicting representations, not sentences

In their preprint, Tommy He of Tufa Labs and Thomas Jiralerspong of Mila, Université de Montréal, study a forecasting problem inside language models. Rather than predict the exact words an answer will contain, they predict a numerical summary of the model’s internal activity during that answer.

The distinction matters. With stochastic decoding, a fixed prompt can produce several different completions. Observing that variation ordinarily requires generating multiple answers. The authors ask whether a separate learned model—a “context–answer metamodel”—can forecast the corresponding answer representations from the activations already available after the language model reads its context.

Their target is a mean-pooled answer activation: a vector formed by averaging hidden states across answer tokens at a selected layer. This compresses an entire answer into one representation. It is not a forecast of every intermediate state or a reconstruction of the generated text.

Four ways to make a forecast

The researchers compare four metamodel families. Linear ridge regression and a multilayer perceptron, or MLP, predict the expected answer vector across repeated generations. A Gaussian model places a probability distribution around the frozen MLP prediction. A conditional flow model learns to transform random noise into samples of possible answer vectors.

The linear predictor receives the activation at the final context token. The main neural configurations instead compress the context activation sequence into 32 bins and combine them through learned attention pooling. This gives them information from across the prompt rather than just its last position.

The primary experiments use Qwen3.5-9B’s decoder block 18. Training draws on 500,000 LMSYS-Chat-1M prompts, with separate validation and test sets of 1,000 and 3,000 prompts. Each prompt has four independently seeded answer generations. Generating these training targets remains necessary: the proposed shortcut concerns forecasting after training, not eliminating answer generation from the research process.

More prompts improve the forecasts

The authors report that linear ridge regression initially outperforms the nonlinear alternatives when training data are scarce. As prompt coverage increases, the neural models overtake it, while the linear model’s performance levels off earlier.

At 500,000 training prompts, the standard 32-bin MLP achieves a mean R² of 0.8739 on held-out LMSYS prompts. This score measures agreement with empirical answer vectors averaged across rollouts; it does not mean that the system predicts 87% of answer words correctly.

The same LMSYS-trained MLP reaches R² values of 0.7309 on WeirdChat and 0.6839 on IFEval without training on those evaluation datasets. A separate representation comparison finds that retaining the full context trace improves the MLP’s results further.

Forecasting variability is less clear-cut. The Gaussian has a slightly better in-domain energy score, a measure of distribution forecast quality, while the flow scores better on the two out-of-domain sets. Neither perfectly matches observed spread: on LMSYS, their generated-to-observed variance ratios are 1.532 and 1.299, respectively, against an ideal value of one.

Behavior signals survive, with qualifications

The researchers also test whether predicted activations retain information associated with behaviors such as safety failures, self-harm-related responses, and deception or false capability claims. They train logistic-regression probes on real answer activations, then apply those probes to forecast means or samples.

This is a separate experiment using newly fitted metamodels trained on both LMSYS and the WeirdChat training split—not the zero-shot evaluation described above. The MLP improves precision–recall performance over direct context-based prediction in all four reported behavior groups, although some improvements are small. Modeling a distribution does not consistently beat using the MLP mean.

Every context-only forecast assigns the same behavior probability to all four answers from a prompt. It therefore estimates prompt-associated risk rather than identifying which particular sampled answer will exhibit a behavior.

Transfer does not remove the limits

Cross-model experiments train metamodels to map Qwen2.5 context activations to Qwen3.5 or Gemma 4 answer representations. Mean forecasts retain much of the accuracy obtained using context activations from the target model itself. These are trained cross-model mappings, not untrained transfer to arbitrary models.

Gemma 4 exposes an important weakness: both distribution models generate more than twice the observed average within-prompt variance, despite accurate mean forecasts.

The preprint also limits its conclusions to selected layers and generation settings. Behavior labels come from GPT-6 Luna judgments without extensive human validation. Forecasts depend on decoding rules, chat templates, generation caps, and pooling choices. The results support further investigation of pre-generation monitoring, but do not demonstrate dependable prevention of harmful answers.