Inference-Time Intervention: Eliciting Truthful Answers from a Language Model

In plain words

Language models often "know" the right answer and still write a false one. This paper looks inside the model while it reads a question and a candidate answer. It finds a small set of internal components, called attention heads, whose signals tell true answers from false ones. Then, while the model writes, it nudges those components a little in the "true" direction at every step. No retraining and no new training text are needed. The directions come from a few hundred labelled examples.

The authors test this on a benchmark of questions where people often answer with a common misconception. The nudge makes the models more truthful, and the gains are largest on an instruction-tuned model. The cost is a trade-off: a stronger nudge makes the model more truthful but less informative, because it answers "I have no comment" more often. The gains are measured on one benchmark, and the nudge strength has to be tuned.

TL;DR

• Find heads whose activations linearly predict truthfulness, keep the top K, and add a scaled truthful direction to each kept head's output at every generated token §3.3 eq. 2.
• Needs white-box access, labelled true and false answers, and a fine-tuned GPT judge for evaluation §4.1 [README.md#L104-L106].
• On TruthfulQA, author-reported true*informative scores: Alpaca 32.5% to 65.1%, Vicuna 51.5% to 74.0% tab. 2.
• Strength alpha trades truthfulness against informativeness; K=48 and alpha=15 were chosen on a sweep for LLaMA-7B §4.3 fig. 4.
• The released code and the paper differ in several places (tuning set for sigma, bias baking, normalisation). Read Paper vs code before implementing edit_weight.py.

Problem

Models can produce fluent false answers even when probes on their hidden states show the truth is represented. On LLaMA-7B and TruthfulQA the authors measure a gap of about 40 points between probe accuracy and generation accuracy §1. Existing fixes such as RLHF and RLAIF need large annotation or compute budgets, and the training can reward pleasing a rater rather than being truthful §1.

Method

The method has two parts. First, a probing step on labelled answers picks a sparse set of attention heads and a direction per head §3.2 §3.3. Second, at inference time the chosen heads' outputs are shifted along those directions eq. 2.

Data. TruthfulQA has 817 questions in 38 categories. It is built so that some humans answer wrongly because of false beliefs or misconceptions §3.1. The authors split it into N = 5,918 question-answer pairs, each labelled 1 if the answer is truthful and 0 if not §3.1. Each pair is formatted as the question followed by the answer §3.2.

Architecture notation. Each attention head h in layer l writes a vector through its own output matrix Q. The paper writes the attention block as a sum over heads, with the activation read after attention and before Q §3.1.

x_{l+1} = x_l + \sum_{h=1}^{H} Q_l^h \, \mathrm{Att}_l^h(P_l^h x_l) eq. 1

Probing. For each head in each layer, the authors fit a linear probe, p(x) = sigmoid(<theta, x>), on the head's activation at the last token of each pair §3.2. The probe is trained on a 4:1 train/validation split. Validation accuracy measures how truth-related each head is §3.2 app. C. The highest single-head accuracy is 83.3%, for the 18th head in the 14th layer, and most of the signal sits in early to middle layers §3.2 fig. 2. Concatenating all heads reaches 84.4% §5.5.

Direction. The paper compares three directions per head §3.3 §4.2:

• the probe weight direction, theta, normalised §3.2;
• the mass mean shift: the vector from the mean of false activations to the mean of true activations §3.3 §4.2;
• the direction found by Contrast-Consistent Search, with its sign set using labels §4.2.

Mass mean shift performs best and is used in all later experiments unless stated otherwise tab. 3 §4.3.

Intervention. Only the top K heads, ranked by validation accuracy, are edited §3.3. For each kept head, sigma is the standard deviation of the activations along the chosen direction. The edit adds alpha times sigma times theta to that head's output before Q. Unselected heads have theta = 0 §3.3:

x_{l+1} = x_l + \sum_{h=1}^{H} Q_l^h \left( \mathrm{Att}_l^h(P_l^h x_l) + \alpha \sigma_l^h \theta_l^h \right) eq. 2

The edit is repeated for every next-token prediction and does not depend on the decoding algorithm §3.3.

Baking into weights. Because the edit is one fixed vector per layer, it can be stored as the output-projection bias of that layer, with Q applied to each head's shift and summed over the intervened heads §5.2:

\mathrm{Bias}_l = \alpha \sum_{h=1}^{H} Q_l^h \left( \sigma_l^h \theta_l^h \right) eq. 3

The authors report close to zero run-time overhead §5.2.

Pseudocode. The fit and the decoding-time edit, following the paper. The normalisation of theta is taken from code, because the paper does not state it for mass mean shift.

# Fit the intervention from labelled QA pairs  [§3.2] [§3.3] [app. C]
# Input:  activations (N, layers, heads, d) at the last token; labels (N,)
# Output: for each kept head, a unit direction and a scale sigma
def fit_iti(activations, labels, questions, num_heads_k, alpha):
    # Split the questions 4:1 into train and validation  [app. C]
    train_idx, val_idx = split_by_question(questions, train_fraction=0.8)
    # One linear probe per head, fit on the training split  [§3.2]
    probe_accuracy = {}
    for layer, head in all_layer_head_pairs():
        probe = fit_logistic_regression(
            activations[train_idx, layer, head], labels[train_idx])
        probe_accuracy[layer, head] = score(
            probe, activations[val_idx, layer, head], labels[val_idx])
    # Keep the K heads with the highest validation accuracy  [§3.3]
    kept_heads = top_k_by_accuracy(probe_accuracy, num_heads_k)
    # Mass mean shift: true mean minus false mean  [§3.3] [tab. 3]
    # Unit-length normalisation is from the code, not the paper
    # [code: validation/validate_2fold.py#L171]
    for layer, head in kept_heads:
        true_mean = mean(activations[labels == 1, layer, head])
        false_mean = mean(activations[labels == 0, layer, head])
        theta[layer, head] = unit_vector(true_mean - false_mean)
        # Sigma: std of the activations projected on theta  [§3.3]
        # The tuning set for this std is not clear; see Paper vs code
        projections = activations_tuning[:, layer, head] @ theta[layer, head]
        sigma[layer, head] = std(projections)
    return theta, sigma, kept_heads

# Shift one head's output during decoding  [eq. 2] [§3.3]
# Input:  head output at the last position, shape (d,)
# Output: shifted head output, same shape
def shift_head_output(head_output, theta, sigma, alpha):
    # Add alpha * sigma * theta; the next step applies Q  [eq. 2]
    return head_output + alpha * sigma * theta

Hyperparameters

Results

Author-reported. The setup differs between tables; see the Source column and the protocol rows above.

Main TruthfulQA results on LLaMA-7B with 5% of questions used for training and validation. Baselines were reproduced by the authors tab. 1 §4.3:

Instruction-tuned models, 2-fold cross-validation tab. 2:

Direction comparison on LLaMA-7B, 2-fold cross-validation, with each direction's own grid-searched alpha tab. 3 §4.3:

Head selection on LLaMA-7B, 2-fold cross-validation tab. 5 §5.5:

Transfer without tuning on other QA and knowledge benchmarks, multiple-choice accuracy, with directions and hyperparameters learned on TruthfulQA tab. 4 §5.3:

The differences are 4.7, 1.5 and 4.45 points, computed from tab. 4. The authors call this "somewhat better" and note that MMLU's question format is closest to TruthfulQA §5.3.

Per-category gains: ITI raises the true*informative score across most subcategories with 10 or more questions, with no single category driving the average §5.1 fig. 5.

Trade-off: the number of "I have no comment" answers rises with alpha, and the authors frame this as truthfulness traded against helpfulness §5.4 fig. 6.

Stability: the authors report that results are quite stable across random seeds. Standard deviations over 5 runs are in tab. 6 app. D.

Qualitative answers for every TruthfulQA question, with and without ITI, are in the appendix app. A. In that 2-fold run, overall truthfulness goes from 31.6% to 54.5% and informativeness from 96.7% to 93.3% app. A.

Paper vs code

Code: https://github.com/likenneth/honest_llama at commit 2c6b2179be7b5aa8f0a171688cf9e01b812ca327, MIT licence. The repository is linked from the abstract and belongs to the first author. The README says the pyvene scripts are an updated implementation with similar results, and that the legacy scripts may suit paper replication better [README.md#L118]. The legacy code is the paper-era version; it is cited below where it differs.

• Tuning set for sigma. Paper: sigma is estimated from the training and validation activations §3.3. Code: sigma is estimated from label-free activations of the tqa_gen_end_q set validate_2fold.py validate_2fold.py. Those prompts end with a random extra question, so the captured position is the last token of that question utils.py utils.py.
• Direction normalisation. Paper does not state normalisation for mass mean shift. Code divides each direction by its norm before scaling utils.py validate_2fold.py.
• Edit positions. Paper writes the edit per token eq. 2 §3.3. Code adds the edit only at the last position of each forward call interveners.py. In the prefill this is the last prompt token; the earlier prompt tokens are not shifted. The legacy code does the same with head_output[:, -1, ...] llama_validate_2fold.py. The repo README says the same [README.md#L39].
• Baked-in bias. Eq. 3 applies Q to the shift eq. 3. The active code sets the bias to the raw displacement and leaves the Q multiplication commented out edit_weight.py. The README warns that baked models and edits on base models are not exactly comparable [README.md#L2].
• Required flag. --use_center_of_mass is off by default, but the intervention loop indexes com_directions unconditionally, which is None without the flag validate_2fold.py validate_2fold.py. The README command includes the flag [README.md#L112]. Without it the script fails as written.
• Probe fit. Paper gives no regularisation §3.2. Code uses scikit-learn logistic regression with default L2 regularisation and max_iter=1000 utils.py.
• Mass mean shift set. Code computes the mean difference over the train and validation questions combined utils.py, which matches app. C's "whole development set" app. C.
• Judge model. Paper says the GPT-judge is a fine-tuned GPT-3-13B model §4.1. The repo notebook fine-tunes davinci-002 finetune_gpt.ipynb finetune_gpt.ipynb. The judge has to be trained with the user's own OpenAI key [README.md#L104-L106].
• MC scoring. Paper: candidate answers are compared by conditional probability §4.1. Code sums the log-probabilities of the answer tokens after dropping the first three, with no length normalisation utils.py.
• KL and CE sample. Paper: a subset of OpenWebText §4.1. Code: 100 samples truncated to 128 tokens; KL is divided by sequence length utils.py utils.py utils.py.
• 5% protocol. Table 1 and the sweep use 5% of questions §4.3. The repo has no script for this split. Its validation script uses 2-fold cross-validation validate_2fold.py.
• Few-shot prompt. Paper gives three different sizes for the few-shot baseline (see Pitfalls). The repo's utilities accept a many-shot prefix, but no script sets it utils.py.

Pitfalls

• Alpha and K are not consistent across the paper. Section 4.3 picks K=48 and alpha=15 §4.3. Table 3 lists alpha=20 as optimal for mass mean shift, and says mass mean shift is used everywhere else tab. 3. Alpha is not stated for Table 1 or Table 2. Code defaults are K=48 and alpha=15 validate_2fold.py. Tune alpha on your own model.
• Alpha is in units of sigma. The edit is alpha times sigma times a unit vector eq. 2. An alpha of 20 means 20 standard deviations along the direction app. B. Because sigma depends on the model's activations, it has to be recomputed for each model.
• Random-direction control. Random directions at alpha=20 leave LLaMA-7B nearly unchanged app. B tab. 3. Keep the random baseline to check that gains come from the learned direction.
• Sigma's data. See Paper vs code. A different tuning set changes the scale of every edit.
• Truthfulness vs informativeness. An answer of "I have no comment" counts as true but not informative, so the true*informative metric penalises it §4.1. Stronger alpha raises truthfulness and lowers informativeness §4.3 §5.4 fig. 6. Report both.
• Judge bias and cost. The GPT-judge slightly overestimates truthfulness and underestimates informativeness in the authors' human check, with no numbers given §4.1. The judge is a fine-tuned OpenAI model, so results are not reproducible without access to it.
• Few-shot baseline sizes conflict. Section 4.2 cites a 50-shot baseline from prior work §4.2. Appendix F says the authors appended 10% of TruthfulQA app. F. Section 4.3 says 5% was used "for a fair comparison" §4.3. Check which one a reproduction should use.
• Protocols are not comparable. Table 1 (5% train, 5% validation) and Tables 2 to 5 (2-fold) use different splits. Baseline+ITI is 43.5% in Table 1 and 42.3% in Table 3 for the same model tab. 1 tab. 3.
• Data needs are stated loosely. The abstract says "few hundred examples" (abstract). Section 2 says "as few as 40" §2. Section 4.3 uses 5% of 817 questions for training (about 41) §4.3. Probes need labelled true and false answers for every head.
• Probe position. Each probe reads only the last token of each pair §3.2. The paper does not examine other positions.
• Cost. Run time is nearly free: one vector added per layer §5.2. Fitting needs a forward pass over all labelled pairs with hooks on every head. Baking the edit into weights is cheaper at inference but changes the prompt tokens too, which the code does not match exactly (see Paper vs code).
• Generalisation is narrow. Gains on the transfer benchmarks are small (4.7, 1.5 and 4.45 points) tab. 4. The authors list testing on wider benchmarks and real chat settings as open work §6.
• Repo replication notes are not the authors' results. The repo has a community replication note on Llama-2 and Llama-3 models with mixed truth and informativeness changes iti_replication_results.md. It is not part of the paper.

Lineage

• Builds on: LLaMA for the base models; the probing and activation-editing idea follows prior interpretability work §2 §3.2. Contrast-Consistent Search, arXiv:2212.03827 (no card), is the main unsupervised comparison §2 §4.2.
• Uses the TruthfulQA benchmark, arXiv:2109.07958 (no card) §3.1.
• Compared against: supervised finetuning and RLHF-style training (no card) §4.2; few-shot prompting §4.2; random and probe-weight directions tab. 3; contrast-consistent search tab. 3.
• Superseded by: not stated on this wiki.