PaperlayerView exact revision on arXiv

Source-linked AI summary

Aligning Language Models Towards Counterfactual Invariance

Basar Demir, Emily Shen, Ksheeraj Sai Vepuri, Sourabh Kulkarni, Harihar Subramanyam, Ankit Jain, Nan Jiang, Chirag Nagpal

arXiv:2610.22942v1cs.CLcs.AIcs.LG

TL;DR

When demographic context is irrelevant, language-model responses should not depend on it, but models must still adapt when it matters. The paper trains on counterfactual prompt pairs with reward comparisons designed to reduce sensitivity to prompt wording. The authors report improved open-ended invariance while retaining most helpfulness gains, though their mitigations do not restore the initializer’s invariance.

  • Problem

    Open-ended evaluation and alignment methods provide limited support for selectively enforcing invariance when demographic information is irrelevant while retaining appropriate adaptation when it matters.

    Source: §1 · pp. 1–2

  • Method

    The authors build labeled counterfactual prompt pairs and use reward-based training objectives that compare responses across arms, including scoring them under the same prompt.

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

  • Results

    The method improves open-ended invariance while retaining most of the helpfulness gain, but does not return the model to its demographic-free initializer’s invariance.

    Source: §5 · p. 9, §6 · p. 10

  • Takeaways & Limitations

    Forced-choice benchmarks can miss failures in open-ended responses, so they do not alone certify invariance in generation.

    Source: §5 · p. 9, §6 · p. 10

    The invariant/dependent labels encode contestable normative judgments, and the authors will release the paired dataset but not the evaluation rubrics or splits.

    Source: §5 · p. 9

Abstract

from arXiv · show

Small variations in prompt characteristics can meaningfully alter the semantics, tone, and contextual assumptions of large language model outputs. Such characteristics can involve group variables such as demographics or other sensitive latent attributes. Sensitivity to such input variables is undesirable when demographic information is causally irrelevant to the task, but might still be relevant for generation of semantically coherent responses. This makes the task of enforcing counterfactual invariance while preserving model utility difficult. To study this problem, we first show that standard model alignment recipes (such as reinforcement learning) increase differences in response quality across demographics, while improving overall response quality, ultimately amplifying what is known as counterfactual bias. To address this, we develop a policy gradient technique involving a group-level reward function that operates on matched pairs of counterfactual prompts. For completeness, we build a dataset of controlled counterfactual prompt pairs and design an evaluation that separately measures cases where demographic information is causally relevant to the answer. Finally, we validate our findings on several fairness benchmarks. We demonstrate that our methods reduce counterfactual differences while preserving overall response quality and that the models continue to adapt appropriately when such sensitive information is relevant.

1 INTRODUCTION

Counterfactual invariance requires responses to ignore input features, such as demographics, when those features are causally irrelevant, while still using context when it matters. The paper targets open-ended generation, where standard alignment can amplify bias and forced-choice evaluations miss important differences.

1 INTRODUCTION

  • When demographic information is causally irrelevant, responses should not depend on it; ignoring relevant context, however, can reduce helpfulness.Over-indexing on context can also produce sycophantic, oblivious, or paternalistic responses.Source: §1 · p. 1
  • Forced-choice benchmarks can miss differences in open-ended responses, including changes in implicit assumptions or tone despite fluent wording.The authors argue that evaluation should assess whether responses remain aligned under non-causal input perturbations, not just whether tokens match.Source: §1 · p. 2
  • The paper proposes training on counterfactual prompt pairs, using reward-based comparison to avoid penalizing harmless paraphrases and a shared prompt to reduce reward-model sensitivity to demographics.For invariant pairs, the method uses GRPO; the dataset pairs prompts differing in race or gender and labels whether that difference matters for the task.Source: §1 · p. 2
  • The authors report that standard RL post-training perpetuates bias on public benchmarks, while their method improves invariance and preserves response quality.They also introduce about 2,500 labeled counterfactual prompt pairs curated from four public corpora.Source: §1 · p. 2

2 RELATED WORK

Prior work studies counterfactual fairness through labels, scores, or detectors, while open-ended evaluation remains difficult. Existing approaches leave gaps in distinguishing harmless rewording from substantive differences and in judging when demographic variation is warranted.

Counterfactual fairness

  • Prior counterfactual-fairness work evaluates predictions through labels, scores, or detector verdicts, while invariance is not always desirable when tasks require treating groups differently.The cited work extends counterfactual fairness from protected attributes to other input attributes and to text and LLMs.Source: §2 · p. 3

Invariance evaluation in forced-choice and open-ended settings

  • Forced-choice benchmarks and free-text scoring can miss or confound differences in open-ended responses, while direct text comparisons may mistake harmless rewording for substantive change.First-Person Fairness uses an LLM rater but applies the same stereotype judgment across tasks without specifying whether demographic differences are warranted.Source: §2 · p. 3

Sensitivity of reward models

  • Reward models can react to demographic cues, and GRPO’s within-prompt normalization does not remove preferences among responses that differ across demographic prompts.Prior work addresses reward-model sensitivity with representation penalties or fairness constraints.Source: §2 · p. 3

Bias mitigation and consistency training

  • Existing bias-reward and consistency methods often judge responses individually or rely on a clean prompt and behavior, potentially penalizing harmless changes.The cited GRPO-style target-distribution approach does not discuss prompt matching or counterfactuals.Source: §2 · p. 3

3 METHODOLOGY

The methodology defines counterfactual invariance for open-ended responses and develops reward-based objectives to compare paired prompts. Its final objective scores responses under the same prompt to reduce the effect of prompt-specific reward-scale shifts.

3.1 PROBLEM FORMULATION

  • Paired prompts share task content and template and differ only in the demographic expression; invariance concerns distributions over task-relevant behavior, not exact text.The behavior representation includes conclusions, recommendations, assumptions, supporting information, and tone.Source: §3.1 · p. 4

3.2 REWARD MODELS AND GRPO

  • GRPO standardizes rewards within each prompt’s response group, but does not compare the two arms, and reward-model preferences may differ across prompts.The cited example reports higher scores for the same response under white and male prompts, especially for responses sampled under those prompts.Source: §3.2 · p. 5

3.3 TRAINING OBJECTIVES

  • The objectives range from matching next-token distributions to symmetrizing rewards or matching their expectations; next-token matching also penalizes harmless surface differences.Reward Symmetrization averages scores for a response and its demographic-swapped version, but averaging does not itself penalize the reward gap.Source: §3.3 · p. 5
  • Reward Expectation Matching aligns the arms’ mean rewards, but cross-prompt mean differences can reflect incompatible reward scales rather than behavior.The Bradley-Terry reward objective is unchanged by a prompt-specific constant, although such a constant changes the difference between the arms’ means.Source: §3.3 · pp. 5–6
  • The expected-reward gradient uses leave-one-out means to estimate the gap and centers each reward to avoid a response influencing its own weight.The resulting correction is added to standardized GRPO rewards; α is fixed across responses and arms within a run.Source: §3.3 · p. 6
  • Transformed Reward Expectation Matching compares responses from both arms under each same prompt, so prompt-specific reward shifts cancel.The comparison yields zero gaps when responses to one arm are distributed as the transformed responses to the other.Source: §3.3 · p. 6

4 EXPERIMENTS AND RESULTS

The experiments compare counterfactual-bias mitigation methods with vanilla GRPO using fairness benchmarks, held-out reward, and open-ended response rubrics. Vanilla GRPO performs worst on key invariance measures and all rubric dimensions, while the proposed methods mitigate bias amplification and Transformed Reward Expectation Matching achieves the highest judge score.

4.1 ENCOURAGING INVARIANCE WITH PROPOSED MITIGATIONS

  • Vanilla GRPO lowers every invariant rubric score despite improving marginal reward, while the proposed methods retain higher scores throughout training.Source: Figure 4 · p. 7

4 EXPERIMENTS AND RESULTS

  • The evaluation uses paired counterfactual prompts labeled invariant, dependent, or agnostic, and assesses fairness benchmarks alongside four open-ended response dimensions.The training starts from Qwen3-4B-Base after demographic-free supervised fine-tuning; evaluation includes Difference Awareness, BBQ, and DiscrimEval.Source: §4 · pp. 6–8

4.1 ENCOURAGING INVARIANCE WITH PROPOSED MITIGATIONS

  • Vanilla GRPO ends with the largest DiscrimEval gap and lowest Difference Awareness accuracy; proposed methods mitigate bias amplification while improving marginal reward.All objectives reach similar held-out reward, and Transformed Reward Expectation Matching achieves the highest judge score.Source: Figure 3 · p. 7, §4.1 · p. 9
  • Vanilla GRPO scores lowest on all four open-ended rubric dimensions, while Transformed Reward Expectation Matching scores highest.Its remaining failures are mostly unsolicited commentary or disclaimers rather than substantive differences.Source: §4.1 · p. 9
  • Transformed Reward Expectation Matching keeps own-arm and transformed opposite-arm reward distributions close, with ∆µ between 0.41 and 0.97.After vanilla GRPO, the distributions separate; within-prompt normalization cannot remove reward shifts that are not shared across responses in a group.Source: Figure 5 · p. 8, §4.1 · p. 9

5 DISCUSSION & LIMITATIONS

The discussion describes trade-offs and limits of counterfactual-bias mitigation under helpfulness alignment. The methods slow bias growth but do not restore the initializer’s invariance, and the labeling framework encodes contestable normative judgments.

5 DISCUSSION & LIMITATIONS

  • Response Matching performs well on forced-choice benchmarks but trails in open-ended generation, where its finite-weight next-token regularizer does not stabilize answer content and register.Source: §5 · p. 9
  • Every mitigation objective slows counterfactual-bias growth under RL but none prevents it or returns the model to the demographic-free initializer’s invariance.The methods retain most of the helpfulness gain; Reward Expectation Matching still ends with a DiscrimEval gap far above the initializer’s.Source: §5 · p. 9
  • The invariant/dependent/agnostic labeling schema encodes contestable normative claims, so the authors release the paired dataset but not its evaluation rubrics or splits.They frame their contribution as the problem formulation and class decomposition, not the alignment target itself.Source: §5 · p. 9

6 CONCLUSION

The paper studies counterfactual invariance in open-ended generation and finds that reward-model scores can shift with demographic prompt wording even for identical responses. Its methods reduce open-ended differences while retaining most helpfulness gains.

6 CONCLUSION

  • Reward Expectation Matching improves open-ended invariance while retaining most of the helpfulness gain from standard RL post-training.The conclusion also reports that standard RL improves helpfulness but amplifies counterfactual bias, partly because the reward model reacts to demographic cues.Source: §6 · p. 10
  • Existing forced-choice benchmarks miss failures that appear only in open-ended responses; future work includes cue-insensitive reward models and reasoning models.Source: §6 · p. 10

A REWARD MODEL BIAS UNDER DEMOGRAPHIC PROMPT SWAPS

  • The reward-model test scores each unchanged response under both demographic prompts, isolating prompt wording as the source of any reward difference.The analysis uses invariant pairs, 16 sampled responses per prompt, and Skywork-Reward-V2-Qwen3-8B.Source: A REWARD MODEL BIAS UNDER DEMOGRAPHIC… · p. 13
  • 0.331 versus 0.153: identical responses sampled under white versus black prompts receive different reward shifts, with a similar pattern for gender.The average shift favors white and male prompts by about 0.24; the own-prompt bonus is about 0.09 for race and 0.15 for gender.Source: A REWARD MODEL BIAS UNDER DEMOGRAPHIC… · p. 13, Table 1 · p. 13

B IMPLEMENTATION OF τ

The paper implements τ as a deterministic lexical swap of supported demographic terms, with rules for pronouns and exceptions for protected spans. Its lexical scope also limits which cues it can change.

B IMPLEMENTATION OF τ

  • τ swaps supported gendered terms and black/white across whole words in one pass, preserving capitalization.The mapping includes nouns, pronouns, and Mr./Ms.; a mixed phrase such as “men and women” becomes “women and men.”Source: B IMPLEMENTATION OF τ · p. 13
  • Context rules resolve ambiguous “her” forms, and applying τ twice restores the original response.For example, “her” becomes “him” after verbs such as “gave” or before words such as “tomorrow,” and otherwise becomes “his.”Source: B IMPLEMENTATION OF τ · p. 13
  • τ leaves quotations, code, URLs, email addresses, and fixed non-person phrases unchanged; its lexical scope misses names and implicit cues.It can also swap terms referring to people other than the focal subject, and rare pronoun constructions may be misassigned.Source: B IMPLEMENTATION OF τ · pp. 13–14

C ANALYSIS OF REWARD SYMMETRIZATION

Reward symmetrization scores responses under both prompts and averages those rewards, but does not itself guarantee equal training rewards or semantic invariance across arms.

C ANALYSIS OF REWARD SYMMETRIZATION

  • Symmetrization rewards responses that score well under both prompts by averaging native and transformed rewards.The method is motivated by the hypothesis that responses scoring well under both prompts serve both demographic contexts better than responses scoring well only under their own prompt.Source: C ANALYSIS OF REWARD SYMMETRIZATION · p. 14
  • A response’s reward gap affects training credit linearly and preserves its sign; the absolute gap is not a penalty.A response scoring 10 and 2 averages 6 and receives more credit than one scoring 5 and 5.Source: C ANALYSIS OF REWARD SYMMETRIZATION · p. 14
  • The training signal changes only when reward gaps vary across responses, and symmetrization neither equates arm rewards nor guarantees semantic invariance.Its effect depends on transformed rewards varying with response content and on τ preserving non-demographic content.Source: C ANALYSIS OF REWARD SYMMETRIZATION · p. 15
  • Symmetrization is unaffected by prompt-specific reward offsets, which are removed by group standardization.Source: C ANALYSIS OF REWARD SYMMETRIZATION · pp. 14–15

D DERIVATIONS FOR REWARD EXPECTATION MATCHING

The derivations turn expected-reward gaps between counterfactual prompt arms into policy-gradient corrections, using leave-one-out estimates and centering. The transformed-reward variant applies the same estimation logic to cross-prompt scores.

D.1 REWARD EXPECTATION MATCHING

  • Reward Expectation Matching targets equal expected rewards across the two prompt arms by adjusting each arm in opposite directions when their means differ.When one arm’s expected reward is higher, the objective favors decreasing it and increasing the other arm’s expected reward.Source: §D.1 · p. 15
  • The policy-gradient correction weights each response by its reward times the expected-reward gap between arms.Because the gap is unobserved, the derivation estimates it with independent-arm samples and a leave-one-out mean to avoid including the response in its own gap estimate.Source: §D.1 · pp. 15–16
  • Centering rewards against the other responses in the same arm leaves the expected gradient unchanged and makes the correction invariant to a constant reward shift.The derivation uses the centered rewards and leave-one-out gap in the resulting corrections.Source: §D.1 · p. 16

D.2 TRANSFORMED REWARD EXPECTATION MATCHING

  • For transformed-reward matching, the derivation groups terms by the arm that sampled each response and applies leave-one-out gaps and centering to each term.A transformed score under the counterpart prompt still uses the score-function factor for the prompt that sampled the response.Source: §D.2 · p. 16

E DATASET CONSTRUCTION

The section builds a relation-labeled dataset of controlled race and gender counterfactual prompt pairs, with a separate demographic-free prompt set. Pairs are generated, reviewed, and split to support studying when demographic context should change a response.

E DATASET CONSTRUCTION

  • The dataset contains 2,582 counterfactual pairs—1,636 race pairs and 946 gender pairs—and 725 unpaired demographic-free prompts.It draws on HealthBench, Anthropic Helpful Base, HelpSteer2, and Wikidata People.Source: E DATASET CONSTRUCTION · p. 16, Table 3 · p. 16
  • Pairs are labeled invariant when responses should preserve advice and quality, dependent when demographic context changes correct content, or agnostic when either response relation is acceptable.The demographic-free class contains prompts without counterfactual pairs.Source: E DATASET CONSTRUCTION · pp. 16–17
  • Pair construction creates black/white or female/male versions differing only in the selected demographic expression and required grammatical agreement.Prompts without an eligible human subject enter the demographic-free class.Source: E DATASET CONSTRUCTION · p. 17
  • A fixed-policy review checks naturalness, answerability, and demographic validity, and keeps variants from each source example within the same split.Valid pairs receive one of the three relation labels.Source: E DATASET CONSTRUCTION · p. 17

F TRAINING CONFIGURATIONS

Training experiments share a demographic-free supervised-fine-tuned initializer and use GRPO post-training. The section selects a reward-matching hyperparameter and describes held-out evaluation folds.

F.1 SUPERVISED FINE-TUNING

  • The shared initializer is Qwen3-4B-Base fine-tuned on 10,000 demographic-free prompts, with the final assistant target alone supervised.The corpus combines 6,702 Anthropic Helpful Base, 2,954 HelpSteer2, and 344 HealthBench prompts, split 90/10 for training and validation.Source: §F.1 · p. 17

F.2 GRPO

  • GRPO samples a response group per prompt and standardizes scalar values using the group mean and standard deviation, with 10^-4 as a numerical stabilizer.Training settings are shared across conditions unless the data or named objective changes.Source: §F.2 · p. 17, Table 5 · p. 18

F.3 HYPERPARAMETER TUNING FOR REWARD EXPECTATION MATCHING

  • α = 0.1 is selected for all experiments because it gives the highest invariant-only judge score while keeping Skywork reward near vanilla GRPO.Under mixed training, the tested values perform similarly on invariant pairs; on dependent pairs α = 0.05 yields the lowest difference-required accuracy while preserving dependent judge score across values.Source: §F.3 · p. 18, Figure 6 · p. 18, Figure 7 · p. 18

G.1 EVALUATION SPLITS

  • Four stratified folds keep counterfactual pairs and shared conversation roots together, with each run training on three folds and evaluating on the held-out fold.Invariant-only models are evaluated on invariant pairs, while mixed-dataset models are evaluated on invariant, dependent, and agnostic pairs.Source: §G.1 · p. 19

G.2 OPEN-ENDED RESPONSE EVALUATION

The evaluation combines open-ended response rubrics, fairness benchmarks, and held-out reward measurements to test both invariance and appropriate demographic adaptation. In mixed training, mitigation improves invariant and agnostic responses while largely preserving dependent response quality, but some benchmark gaps remain.

G.2 OPEN-ENDED RESPONSE EVALUATION

  • Open-ended evaluation samples four responses per arm and uses three LLM judges; invariant and dependent rubrics respectively assess consistency under demographic swaps and appropriate use of relevant demographic context.Judges score criteria from 0 to 4, and median scores are scaled to 0–100.Source: §G.2 · p. 19, §G.3 · pp. 19–20

G.4 INTER-JUDGE AGREEMENT

  • All three judges agree exactly on 65.2% to 81.0% of ratings, with mean pairwise κw between 0.260 and 0.386; the authors characterize agreement as fair and ratings as noisy.Comparisons average median scores over many pairs and four responses per arm to reduce noise.Source: §G.4 · p. 20, Table 6 · p. 20

G.5 FAIRNESS BENCHMARKS

  • Fairness evaluation uses Difference Awareness, BBQ, and DiscrimEval alongside held-out Skywork reward on demographic-free prompts, restricting benchmark groups to black/white and female/male comparisons.Difference Awareness separates equal-treatment from difference-required items; BBQ measures disambiguated accuracy and ambiguous bias, while DiscrimEval measures demographic effects on Yes/No decisions.Source: §G.5 · p. 20, §G.6 · p. 21

H.1 INVARIANT PAIRS

  • Both mitigation objectives achieve held-out Skywork reward close to vanilla GRPO and outperform it on substance invariance, tonal neutrality, and unsolicited commentary; expectation matching gives the largest gains on the first two.On equality-required Difference Awareness items, expectation matching retains early accuracy while vanilla GRPO and Reward Symmetrization lose accuracy.Source: Figure 9 · p. 21, §H.1 · p. 21, Figure 10 · p. 22

H.2 DEPENDENT PAIRS

  • Both mitigation objectives keep dependent judge scores near the initializer as vanilla GRPO declines, while utilization remains similar across methods and above the initializer.Vanilla GRPO and Reward Symmetrization gain more accuracy than expectation matching on difference-required Difference Awareness items; Figure 12 attributes vanilla GRPO’s judge-score drop to judiciousness and tonal appropriateness.Source: Figure 11 · p. 22, Figure 12 · p. 22, §H.2 · p. 22

H.3 AGNOSTIC PAIRS

  • On agnostic pairs, invariant-mitigation checkpoints score higher than vanilla GRPO under both invariant and dependent rubrics, regardless of which rubric is applied.Source: Figure 13 · p. 23, §H.3 · p. 23

H.4 BBQ

  • On BBQ, disambiguated accuracy stays above the initializer and ambiguous bias falls below it early for all methods; bias later grows under vanilla GRPO but remains lower under both mitigation objectives.Expectation matching has the lowest ambiguous bias at the end of training, while its label-prompt accuracy drops at the last checkpoint.Source: Figure 14 · p. 23, §H.4 · p. 23

H.5 SUMMARY

  • Mitigating only invariant pairs improves invariant rubric scores, keeps dependent response quality near the initializer, preserves demographic use, and raises agnostic-pair scores, though the DiscrimEval gap remains and expectation matching lowers difference-required accuracy.The section concludes that the selective objective does not close every gap.Source: §H.5 · p. 24

I.1 BBQ UNDER INVARIANT-ONLY TRAINING

Invariant-only training produces small differences between methods on BBQ, despite variation across checkpoints. Across objectives, mitigation lowers the reward gap between own-arm and transformed opposite-arm responses.

I.1 BBQ UNDER INVARIANT-ONLY TRAINING

  • Expectation matching ends with the highest label-prompt accuracy and lowest ambiguous bias, but the differences between methods are small relative to checkpoint variation.Source: §I.1 · p. 24
  • Every invariant-only run exceeds the initializer in disambiguated BBQ accuracy and stays below it in ambiguous bias for label and name prompts.The methods show small differences relative to checkpoint variation.Source: Figure 15 · p. 24, §I.1 · p. 24

I.2 FULL REWARD DISTRIBUTIONS

  • Every mitigation objective lowers ∆µ below vanilla GRPO for all four groups; Reward Symmetrization and expectation matching keep it lowest.∆µ is the difference between mean rewards for own-arm and transformed opposite-arm responses.Source: Figure 16 · p. 25, §I.2 · p. 25

J SAMPLE GENERATIONS

The appendix presents paired generations from trained checkpoints, using highlighting and pair scores to make demographic-dependent differences in open-ended responses visible.

J SAMPLE GENERATIONS

  • Paired generations are shown side by side; yellow marks demographic-dependent passages, pair scores use the lower arm’s mean judge score, and lower Output JSD and policy JSD indicate more similar token-level responses.The sample tables cover Response Matching, Reward Symmetrization, and Transformed Reward Expectation Matching alongside vanilla GRPO at step 1,000.Source: J SAMPLE GENERATIONS · p. 25, Table 7 · p. 26, Table 8 · p. 29
  • The change-management examples range from a gendered account of a female executive’s value to descriptions emphasizing skills and responsibilities regardless of gender.Other samples include contrasting ride recommendations and paired financial advice.Source: J SAMPLE GENERATIONS · pp. 26, 29
  • The savings examples offer extensive advice to both a white and a Black user, but frame the guidance and supporting claims differently across the two prompts.The white-user response says race has no bearing on savings goals, while the Black-user response adds Black-specific resources and statistics.Source: J SAMPLE GENERATIONS · p. 27
  • The mortgage examples give both users fixed-rate and variable-rate comparisons, while the woman-directed response adds gender-specific claims about preferences and financial risk.The man-directed response instead presents general pros, cons, and choice considerations.Source: J SAMPLE GENERATIONS · p. 30

Related papers

Loading 2610.22942v1…