Counterfactual Token Generation in Large Language Models
Ivi Chatzi, Nina Corvelo Benz, Eleni Straitouri, Stratis Tsirtsis, Manuel Gomez-Rodriguez
cs.LG, cs.AI, cs.CL
2024-09-25
A Gumbel-Max SCM over the LLM sampler, replaying RNG states, yields counterfactual continuations at near-zero extra cost on Llama 3 8B and Ministral-8B.
“What would the rest of the text have been if that word had been different?” People use that kind of counterfactual to learn from sparse experience. Serving LLMs cannot, and the reason sits in the sampler. Given a prompt, the network maps context to a token distribution and a sampler draws. The process is stateless: the sampler does not record how that particular randomness would have behaved after an intervention. Replacing a generated “blue” with “purple” and continuing is a fresh interventional draw. The noise does not match, and neither does the later structure. Reusing a random seed is not enough, because the input length changed and the sampler state is already off.
MPI-SWS is not trying to match human counterfactuals (that literature usually fine-tunes). The goal is counterfactuals consistent with this LLM’s own world model, at nearly the cost of ordinary generation.
Autoregression is written as a structural causal model. The network produces a token distribution at each step; the next token is fT(distribution, exogenous noise Ui). To continue counterfactually, intervene on some tokens and reuse the original Ui on later steps. Without assumptions on fT and the noise law, the counterfactual is not identified: many pairs (fT, PU) match the observed distribution and disagree on “what if.”
They pick Gumbel-Max: token = argmaxt (log dt + Gumbelt). This SCM satisfies counterfactual stability: unless the intervention lowers that token’s relative chance against the others, a switch is unlikely. Later text tends to reuse tokens from the original. Implementation does not store Gumbel vectors (O(vocab × length) floats). It stores the RNG state per step and replays it, so extra memory is O(length) integers. No fine-tuning, no prompt change.
Top-k and nucleus sampling can restrict the argmax to a candidate subset, and then stability is no longer guaranteed. Temperature still sharpens or flattens the distribution.
The model is implemented on Llama 3 8B-Instruct and Ministral-8B-Instruct.
Qualitative story: a fantasy captain prompt produces Lyra. Replace the name with Maeve in the first sentence and continue. Interventional generation (fresh noise) diverges after a few tokens. Counterfactual generation matches the original almost token-for-token at first, then still forks, so a name change already moves later distributions. Changing the ship name, dropping “trusty,” or replacing “sea” with “blue” forks as well.
Quantitative: 2,000 questions from LMSYS Chat 1M, two random token replacements per output (one in each half), 4,000 sequences. Edit distance is measured on the regenerated suffix. Counterfactual generation sits below interventional generation across temperature, top-p, and top-k. Llama hugs the original more than Ministral. Top-p and top-k variants have no stability proof and still hug more in practice.
Bias detection is the demo. Both models generate fictional census records (114 people for Llama, 158 for Ministral). Sex and race are then intervened. On Llama, most male incomes stay or fall if the person had been female, and most female incomes stay or rise if male; the total effect has larger variance than the direct effect. Male and female factual income share the same median, so the effect is invisible in the observational histogram. On Ministral, education rises for Native and Hawaiian individuals under other races and falls for Asian Americans; Asian occupations shift from STEM toward humanities if the person had been Black.
Adding counterfactuals to an open model is a small change: swap in Gumbel-Max and keep RNG states. Three uses follow. Probe causal dependence inside the model, as in the census demo. Quantify which tokens in a completion carry a later claim, in the spirit of feature attribution. In an editor, a user changes a few words and wants the rest both updated and similar. What is being measured is the model’s world, not ours.
Gumbel-Max is one SCM with stability. A different mechanism yields a different family of counterfactuals. Haugh and Singal have argued that stability needs a domain justification and can admit counterfactuals it was meant to block. The appendix reruns with a classical categorical sampler that does not satisfy stability. When model-world counterfactuals disagree with human causal judgment, they are not explanations. Only two 8B instruct models are tested; how sensitivity scales with size is unknown. The census study is a bit over a hundred fictional records. Those associations live in the model’s world and do not transfer to society.