PaperlayerView exact revision on arXiv

Source-linked AI summary

Counterfactual Test-Time Augmentation for Shortcut Removal in Language Models

Katie Matton, Rosalind W. Picard, John Guttag, Maggie Makar

arXiv:2610.33904v1cs.LGcs.AIcs.CL

TL;DR

LLMs can rely on training-data shortcuts that may fail when correlations shift, while many existing removal methods require resources users may lack. The paper proposes counterfactual test-time augmentation and an adaptive selection procedure; it reports that Select-CTTA outperforms test-time baselines and often approaches balanced-data fine-tuning. Its theory assumes known shortcuts and depends on distribution-shift and editor-accuracy conditions that may fail in practice.

  • Problem

    Many shortcut-removal methods operate at training time and require labeled data, model-weight access, and substantial compute, while shortcut correlations may not hold at test time.

    Source: §1 · p. 1

  • Method

    CTTA averages LLM predictions across auxiliary-model edits that change shortcut values, and Select-CTTA uses diagnostic tests to select among CTTA weighting methods.

    Source: §1 · pp. 1–2, §3 · p. 3

  • Results

    Select-CTTA outperforms test-time baselines and often approaches shortcut-balanced fine-tuning; across two text-classification benchmarks and six LLMs, it achieves the best or near-best performance among CTTA methods.

    Source: §1 · p. 2

  • Takeaways & Limitations

    CTTA can substantially reduce shortcut reliance at test time, and Select-CTTA adaptively selects among methods without requiring labeled test data.

    Source: §1 · p. 2, §6 · p. 10

    CTTA assumes the shortcut and its possible values are known, and its theory relies on distribution-shift and counterfactual-editor accuracy assumptions that may fail in practice.

    Source: §6 · p. 10

Abstract

from arXiv · show

Large language models (LLMs) can exploit shortcuts, i.e., correlations present in training data that may not hold at test time. This can lead to poor predictions when correlations shift. Most existing solutions for removing shortcuts operate at training time, requiring labeled training data, model weight access, and significant computational resources. Here, we study the problem of test-time shortcut removal. We propose a family of counterfactual test-time augmentation (CTTA) methods: given an example, each method uses an auxiliary model to generate counterfactual edits in which the shortcut value is changed, and then takes a weighted average of the LLM's predictions across these edits. We prove that the optimal shortcut-invariant predictor is a member of this CTTA family, and we characterize its weights. The weights depend on the training distribution, and therefore are, in general, not identifiable from the trained model and unlabeled test data. However, we present assumptions under which the weights can be obtained using the token probabilities of the trained LLM, and propose tests to detect violations of these assumptions. We use these tests to develop Select-CTTA, an aggregate method that selects among CTTA methods. Empirically, we find that Select-CTTA outperforms test-time baselines and approaches the performance of models fine-tuned on shortcut-balanced data.

1 Introduction

The paper addresses shortcut reliance without requiring the labeled data, model-weight access, and substantial compute often needed by training-time methods. It proposes CTTA, analyzes when its predictors are optimal, and introduces Select-CTTA to choose among methods; experiments find it outperforms test-time baselines and often approaches balanced-data training.

1 Introduction

  • Most existing shortcut-removal methods operate at training time and require labeled data, model-weight access, and substantial compute that users may lack.The paper therefore studies shortcut removal at test time without model-weight access or labeled data.Source: §1 · p. 1
  • CTTA changes the shortcut value in counterfactual edits and averages the LLM’s predictions across those edits.Its predictors are shortcut-invariant when their weights do not depend on the shortcut value in the example.Source: §1 · pp. 1–2
  • Under a common shortcut-learning model, the optimal shortcut-invariant predictor belongs to the CTTA family, but its weights generally cannot be identified without training-distribution information.The paper gives assumptions under which observable token probabilities can recover the weights, yielding Model-Weighted CTTA.Source: §1 · p. 2
  • Select-CTTA uses diagnostic tests to choose Model-Weighted CTTA when checks pass and uniform weights otherwise.The fixed-weight fallback avoids weight estimation; fixed weighting achieves the optimal shortcut-invariant predictor’s AUROC, though it is miscalibrated.Source: §1 · p. 2
  • Across two text-classification benchmarks and six LLMs, Select-CTTA consistently achieves the best or near-best performance among CTTA methods and outperforms test-time calibration baselines.In many cases, it approaches models trained on gold-standard balanced data while requiring substantially fewer resources.Source: §1 · p. 2

2 Problem Setting

The problem setting considers test-time prediction when the association between a shortcut and the label may change. It defines shortcut-invariant risk and specifies the data-generating assumptions and resources available to the predictor.

2 Problem Setting

  • The assumed causal process has label Y affect text X only through label-relevant content X_Y, while shortcut S may correlate with Y.The shortcut and its possible values are assumed known, but the functions recovering X_Y and S from text are not assumed known.Source: §2 · p. 2, Figure 1 · p. 3
  • The test distribution may change the association between shortcut S and label Y while keeping the other specified distribution components fixed.The test conditional P(S | Y) is unrestricted, subject to stated positivity conditions in the training distribution.Source: §2 · p. 2
  • A predictor is shortcut-invariant when its risk is unchanged across distributions in the specified family, regardless of the shortcut–label relationship.The risk is defined under logistic loss.Source: §2 · p. 3
  • The optimal shortcut-invariant predictor is the member of the invariant set with the lowest risk and, under the assumed process, equals the conditional label distribution given label-relevant text content.The cited passage attributes this characterization to Makar et al. (2022).Source: §2 · p. 3
  • The assumed test-time resources include a trained LLM queryable in classification and likelihood modes, a counterfactual shortcut editor, and unlabeled test texts.The LLM modes are assumed to recover the training conditional and training text marginal, respectively; the editor is assumed to preserve label-relevant content and produce edits in training support.Source: §2 · p. 3

3 Methods

CTTA averages LLM predictions on counterfactual edits, with weights controlling shortcut invariance and optimality. The methods develop model-based and fixed weighting schemes, then use diagnostics to select between them.

3.1 The CTTA Family

  • CTTA evaluates edits that vary the shortcut and averages their predictions; its optimal shortcut-invariant weights marginalize the shortcut conditional on label-relevant content.The optimal predictor belongs to the CTTA family, but its weights generally depend on the training distribution.Source: §3.1 · p. 4

3.2 Model-Weighted CTTA

  • Under likelihood-classifier alignment and an idealized editor, Model-Weighted CTTA recovers the optimal shortcut-invariant predictor.The editor assumption requires residual text variation across edits to contribute equally to their training probabilities.Source: §3.2 · pp. 4–5

3.3 Fixed-Weight CTTA

  • Fixed-Weight CTTA avoids weight estimation; Uniform CTTA distributes weight equally, while Neutral CTTA assigns it all to a designated neutral value.Any fixed weights preserve the optimal predictor’s ranking and achieve the same AUROC, although fixed-weight predictors may be miscalibrated.Source: §3.3 · p. 5

3.4 Select-CTTA

  • Select-CTTA uses unlabeled test-batch diagnostics to choose Model-Weighted CTTA when checks pass and Uniform CTTA otherwise.Its checks assess shortcut strength, weight collapse, and likelihood-classifier alignment; the collapse check alone is not comprehensive.Source: §3.4 · pp. 5–6

4 Experiments

The experiments test CTTA on math-solution correctness and occupation classification under reversed shortcut-label associations, across fine-tuned and zero-shot LLMs. Select-CTTA performs strongly across these settings, and its diagnostics generally choose the better-performing CTTA method.

4.1 Tasks

  • The benchmarks are math-solution correctness with inserted opinion statements and attorney-versus-paralegal classification with gender as the shortcut.The math task uses PRM800K examples; occupation classification uses Bias in Bios biographies.Source: §4.1 · p. 7

4.2 Experimental Settings

  • Experiments reverse the shortcut-label association at test time and cover six instruction-tuned LLMs under all-tokens fine-tuning, completion-only fine-tuning, and zero-shot evaluation.The main controlled-shortcut experiments use p = 0.95; the test label prevalence is 0.5.Source: §4.2 · p. 7
  • Comparisons include CTTA variants, test-time calibration baselines, and balanced-data fine-tuning as a reference; ranking and calibration are measured with AUROC and ECE.Balanced-data fine-tuning is a reference because it assumes labeled training data and model-weight access.Source: §4.2 · p. 7

4.3 Results

  • Select-CTTA achieves the best overall performance among CTTA methods, tracks the stronger method, and remains competitive with balanced-data models without labeled data or weight access at test time.CTTA methods substantially outperform test-time baselines in most fine-tuned settings.Source: Table 1 · p. 8, §4.3 · p. 8
  • Mean AUROC rises from 0.113 for the unaugmented model to above 0.60 with CTTA, versus 0.552 for the strongest existing baseline on zero-shot math.Across both tasks, zero-shot CTTA variants are generally competitive with or outperform existing test-time baselines.Source: Table 2 · p. 8, §4.3 · pp. 8–9
  • In fine-tuned models, Model-Weighted CTTA performs best with all-tokens training, while Uniform outperforms it with completion-only training, consistent with the stated assumption about jointly optimized input and label tokens.Uniform and Neutral CTTA also show strong AUROC, consistent with fixed-weight CTTA's optimal-ranking result.Source: §4.3 · p. 8
  • Select-CTTA selects the better-performing CTTA method most of the time, favoring Model-Weighted for all-tokens fine-tuning and Uniform for completion-only and zero-shot settings.Its label-free diagnostics do not observe the fine-tuning regime; the selection pattern supports using them to assess whether Model-Weighted assumptions hold.Source: Figure 2 · p. 9, §4.3 · p. 9
  • Select-CTTA consistently chooses Uniform in zero-shot experiments, balancing AUROC and ECE across both tasks; on math, Model-Weighted has slightly higher AUROC but substantially worse calibration.The CTTA variants perform similarly on occupation.Source: §4.3 · p. 9

5 Related Work

Prior shortcut-removal work mainly modifies training data or training objectives, while test-time approaches address specific biases or use calibration. CTTA is closest to Out-of-Context prompting, but this work analyzes optimal aggregation weights and recovers them without labels.

5 Related Work

6 Discussion

The paper presents CTTA as a test-time way to reduce shortcut reliance without the resources often needed by training-time methods. It characterizes predictor optimality, evaluates CTTA, and proposes adaptive selection, while identifying assumptions and access requirements that constrain its scope.

6 Discussion

  • CTTA aggregates predictions across counterfactual shortcut edits to reduce shortcut reliance at test time, and the paper characterizes optimality conditions and proposes Select-CTTA to choose among predictors.The discussion frames shortcut reliance as an obstacle to LLM safety and reliability.Source: §6 · p. 10
  • CTTA assumes the shortcut and its possible values are known, while the theory also depends on distribution-shift and counterfactual-editor accuracy assumptions that may fail in practice.The authors identify shortcut discovery and greater robustness to imperfect editors and broader shifts as future-work directions; Model-Weighted CTTA also requires token likelihoods.Source: §6 · p. 10

A CTTA Proofs

The appendix proofs establish when CTTA predictors are shortcut-invariant or optimal, then derive a single-edit log-odds relation and show that fixed weights preserve optimal rankings.

A.1 Shortcut Invariance of CTTA Predictors

  • If each CTTA weight depends on text only through its label-relevant content, the resulting predictor is shortcut-invariant.The proof uses editor assumptions to preserve label-relevant content while setting the shortcut value.Source: §A.1 · p. 13

A.2 Optimality of CTTA

  • For an edited input, the editor preserves label-relevant content and sets the shortcut; conditional independence and training-optimality connect its prediction to the optimal predictor.The proof invokes the law of total probability conditioned on label-relevant content.Source: §A.2 · p. 13

A.3 Optimality of Model-Weighted CTTA

  • Under assumptions A1 and A2, Model-Weighted CTTA recovers the optimal weights.The proof equates its weights with the optimal weights for an arbitrary input.Source: §A.3 · p. 14

A.4 Single Edit Log-Odds Decomposition

  • The single-edit log-odds decomposition expresses the edit's log-odds as the optimal predictor's logit plus λ(s), where λ(s) = log Ptrain(S = s | Y = 1).Source: §A.4 · p. 14

A.5 Rank Optimality of Fixed-Weight CTTA

  • Fixed-weight CTTA induces the same ranking as the optimal predictor and therefore achieves the same AUROC under every distribution in the stated family.The proof shows the fixed-weight predictor is a strictly increasing transformation of the optimal predictor.Source: §A.5 · p. 15

B Select-CTTA Details

Select-CTTA uses unlabeled test-batch diagnostics to choose between Model-Weighted and Uniform CTTA. It uses Model-Weighted only when all checks pass.

B Select-CTTA Details

  • Select-CTTA computes each test from the batch’s classifier log-odds and Model-Weighted CTTA weights for every example and shortcut value.It makes one choice per trained LLM and test batch, then applies that choice to every example.Source: B Select-CTTA Details · pp. 15–16
  • The shortcut-strength test fails when the most different shortcut edits shift mean log-odds by less than δstr.If even the most different edits barely change predictions, estimating weights offers little benefit.Source: B Select-CTTA Details · p. 15
  • The weight-collapse test fails when median weights concentrate too heavily on one edit, either within the principal pair or across all shortcut values.Medians prevent a few examples with legitimately concentrated optimal weights from triggering the test.Source: B Select-CTTA Details · p. 16
  • The alignment test fails when the Spearman correlation between relative weights and edit log-odds falls below δalign.A positive correlation means the weights emphasize s+ more for examples whose content looks more like class 1.Source: B Select-CTTA Details · pp. 15–16
  • The threshold choices trade weight estimation against acceptable log-odds deviation, weight concentration, and ranking agreement.The authors recommend conservative thresholds because inaccurate estimated weights can hurt ranking, while Uniform CTTA retains the optimal ranking under their assumed setting.Source: B Select-CTTA Details · p. 16
  • Model-Weighted CTTA is selected only if every test passes; otherwise Select-CTTA uses Uniform CTTA.The stated thresholds are δstr = 0.5, δcol = 0.8, and δalign = 0.5, fixed before testing.Source: Table 3 · p. 16, B Select-CTTA Details · p. 16

C Dataset Details

The experiments use math-solution grading and occupation classification, with solution opinions and gender serving as shortcuts, respectively.

C.1 Math Grading (Sycophancy)

  • Math examples pair MATH problems with PRM800K step labels, retaining solutions only when human step ratings and normalized final answers jointly support a correct or incorrect label.The study samples 3,200 eligible problems and includes one correct and one incorrect solution per problem, with a 768-token combined-text ceiling.Source: §C.1 · p. 17
  • The math shortcut is an appended opinion sentence with three values, and its label correlation is controlled across training and testing.The opinions are “I think this solution is correct,” “I think this solution is incorrect,” and “I’m unsure whether this solution is correct.”Source: §C.1 · p. 17

C.2 Occupation Classification (Gender)

  • Occupation examples come from attorney and paralegal biographies, use gender as the shortcut, and label paralegal as class 1.Biographies containing explicit profession terms are edited to remove them because those terms make the task trivial.Source: §C.2 · p. 17
  • Because source biographies cannot supply the desired training proportions, counterfactual gender edits expand each occupation class to 1,100 examples.The test set has 291 biographies per class, with edited examples used to create the desired gender distributions.Source: §C.2 · p. 17

D.1 Training Distributions

The appendices specify the controlled shortcut distributions, evaluation setup, baselines, editing pipeline, metrics, and released assets.

D.1 Training Distributions

  • Training shortcut strength p controls the designated shortcut’s label-conditional probability, while balanced reference models use every shortcut value equally across labels.For math the designated value is the incorrect opinion; for occupation it is man.Source: §D.1 · p. 18

D.2 Test Distributions

  • Tests use matched, reversed, and balanced shortcut distributions, with each split reusing the same underlying test examples.At p = 0.95, each split contains 2,600 solutions or 582 biographies.Source: §D.2 · p. 18, Table 4 · p. 18

D.3 Models

  • The evaluation uses six Qwen, Gemma, and Mistral models prompted to answer classification tasks with a single letter.Classification probabilities come from the two answer-letter logits; likelihood scores omit tokens shared across edits.Source: §D.3 · pp. 18–19

D.4 Model Training Settings

  • Fine-tuned models train all weights for six epochs, keep a validation-selected epoch, and use either all-token or answer-only cross-entropy.Untuned models are scored under both class-to-letter assignments, whose probabilities and text likelihoods are averaged.Source: §D.4 · p. 19

D.5 Baselines

  • Baselines include batch and contextual calibration, out-of-context prompting, and the same base model fine-tuned on balanced data.Calibration estimates class priors from unlabeled batches or content-free inputs, while the balanced reference predicts on unaugmented text.Source: §D.5 · pp. 19–20

D.6 Counterfactual Editor

  • The counterfactual editor changes or removes math opinions and flips, neutralizes, or scrubs occupation-text gender and profession signals.Occupation rewrites are judged and checked automatically, and failed rewrites are discarded.Source: §D.6 · p. 20

D.7 Metrics and Confidence Intervals

  • Reported metrics are accuracy, AUROC, cross-entropy, and 15-bin expected calibration error, with 95% percentile bootstrap intervals over 2,000 resamples.Method comparisons use paired resamples.Source: §D.7 · p. 20

D.8 Assets and licenses.

  • The paper uses MIT-licensed PRM800K and Bias in Bios data, Apache 2.0 Qwen and Mistral models, and Gemma 3 under its Terms of Use.Source: §D.8 · p. 20

E.1 Results Across Test Distributions

Across reversed and balanced test splits, Select-CTTA performs among the strongest methods, while methods that retain the useful shortcut perform best on matched data. The appendices also describe results across shortcut strengths and report accuracy and cross-entropy for the reversed split.

E.1 Results Across Test Distributions

  • On reversed and balanced splits, Select-CTTA is consistently among the top-performing methods and balances AUROC and ECE; on matched data, methods that do not remove the useful shortcut perform best.Table 5 reports fine-tuned models at p = 0.95, averaging results over six models.Source: Table 5 · p. 21, §E.1 · p. 21

E.2 Shortcut Strength

  • Table 7 reports results at shortcut strengths p ∈ {0.50, 0.75, 0.95}, each evaluated on its own reversed split.At p = 0.50, the reversed and matched splits coincide.Source: §E.2 · p. 21

E.3 Accuracy and Cross-Entropy

  • Tables 8 and 9 report threshold-0.5 accuracy and cross-entropy on the p = 0.95 reversed split for fine-tuned and zero-shot models.Source: §E.3 · p. 21

E.4 Discussion of Balanced Training Comparison

Select-CTTA balances AUROC and ECE on reversed and balanced splits, while methods that retain shortcut information perform best on the matched split. Model-Weighted CTTA sometimes outperforms shortcut-balanced training in the all-tokens setting, which the authors attribute to complementary label signals.

E.4 Discussion of Balanced Training Comparison

  • In the all-tokens setting, Model-Weighted CTTA sometimes outperforms shortcut-balanced training, which the authors attribute to complementary classification and likelihood-based label signals.The likelihood-based weights provide an indirect label prediction, and their errors may differ from those of the classification predictions.Source: §E.4 · pp. 21–22
  • On the matched split, methods that do not remove the useful shortcut perform best; on the other splits, Select-CTTA achieves the best balance of AUROC and ECE.Source: Table 6 · p. 22
  • As shortcut strength increases, standard predictors and calibration baselines degrade sharply in most cases, while CTTA largely preserves AUROC and calibration.Source: Table 7 · p. 23
  • Under completion-only loss, Model-Weighted CTTA degrades substantially, while Select-CTTA avoids its failures by choosing Uniform CTTA.Source: Table 8 · p. 23
  • On reversed-split zero-shot models, Select-CTTA consistently chooses Uniform CTTA and obtains top performance across all metrics and tasks.Source: Table 9 · p. 24

Related papers

Loading 2610.33904v1…