cs.CLOct 16, 2026
SaveAligning Language Models Towards Counterfactual Invariance
Organizations: Meta Superintelligence Labs · University of California, San Diego · Independent Researcher
Abstract
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.
Figures & tables
Figure 1: Sensitive context such as demographics can subtly affect the quality of the open-ended response. Such dependence might be relevant in certain instances, but actively harmful in other situations and lead to diminished response quality or stereotyping. This paper studies the amplification of such harms as a side effect of training language models with reinforcement learning.
Figure 2: Overview of the proposed objectives for a pair of invariant prompts : The central idea in our paper involves enforcing invariance by encouraging the reward distributions across counterfactuals to be similar. Towards this goal, a natural choice involves matching the first moment , that is the expectation of the rewards.
Figure 3: Bias and marginal reward across methods: Across DiscrimEval and Difference Awareness, we find that helpfulness-only RL with vanilla GRPO amplifies counterfactual bias. Our proposed methods mitigate this amplification while improving marginal reward.
Figure 4: Invariant rubric scores of the proposed training methods versus vanilla RL: We find that although vanilla GRPO improves marginal reward, it lowers every rubric score. Our proposed methods retain higher scores throughout training.
Figure 5: Reward distributions under the same prompt: For each group, we score two sets of responses under that group’s prompt, the responses generated for the group and the responses generated for its counterpart after swapping the demographic references. Vanilla GRPO separates the distributions relative to the initializer, whereas Transformed Reward Expectation Matching keeps them aligned.
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
| Mean by sampling arm | |||||
|---|---|---|---|---|---|
| Axis | Mean | Mean | Favored rate | White/male | Black/female |
| Race | 69.0% | ||||
| Gender | 73.8% | ||||
Table 1: Reward differences for identical responses under demographic prompt swaps. The favored rate is the fraction of comparisons in which the white or male prompt receives the higher reward. The last two columns split responses by the prompt under which they were sampled. Rewards are on the native reward-model scale.
| Input | Output of |
|---|---|
| The female patient should schedule a follow-up. She can call her doctor tomorrow. | The male patient should schedule a follow-up. He can call his doctor tomorrow. |
| I gave her the report, and her manager approved it. | I gave him the report, and his manager approved it. |
| His advice to his son was to trust himself. The idea was his alone. | Her advice to her daughter was to trust herself. The idea was hers alone. |
| As a Black entrepreneur, he should compare loan offers before Black Friday. | As a White entrepreneur, he should compare loan offers before Black Friday. |
| He wrote “my father was a man of few words” in the email to Mr. Smith. | She wrote “my father was a man of few words” in the email to Ms. Smith. |
Table 2: Examples of applied to responses.
| Source family | Invariant | Dependent | Agnostic | Demographic-free |
|---|---|---|---|---|
| Anthropic Helpful Base | 575 | 161 | 213 | 377 |
| HealthBench | 881 | 87 | 176 | 38 |
| HelpSteer2 | 281 | 49 | 77 | 310 |
| Wikidata People | 0 | 82 | 0 | 0 |
| Total | 1,737 | 379 | 466 | 725 |
Table 3: Composition of the dataset by source family and relation class. The counterfactual pairs comprise 1,636 race pairs and 946 gender pairs.
| Class | Desired behavior | Example |
|---|---|---|
| Invariant | Preserve relevant content and tone. | Cooking instructions for a black/white user. |
| Dependent | Use demographic context when relevant. | Notable female/male filmmakers. |
| Agnostic | Accept either valid strategy. | Advice before a blind date. |
| Demographic-free | Respond correctly and helpfully. | What is 5G? |
Table 4: Desired behavior across counterfactual-pair relations and the demographic-free class.
| Training element | Setting |
|---|---|
| Optimizer | Adafactor ( Shazeer and Stern, 2018 ) with relative-step scheduling disabled, learning rate , 30 warmup updates followed by a constant rate, zero weight decay, and gradient clipping at 1.0 |
| GRPO | KL coefficient , clipping coefficient , and rewards standardized within each prompt group |
| Group size | responses per prompt |
| Training decoding | Temperature 1.0, top- , and top- |
| Compute | Each training run takes about 30 hours per fold on 3 GPUs with 96 GB of memory each |
Table 5: Training settings shared across conditions. A condition changes only its training data or explicitly named objective.
Figure 6: Effect of on Transformed Reward Expectation Matching under invariant-only training.
Figure 7: Effect of on Transformed Reward Expectation Matching under mixed-dataset training, evaluated on invariant pairs.
Figure 8: Effect of on Transformed Reward Expectation Matching under mixed-dataset training, evaluated on dependent pairs.
| Data and rubric | All three exact | All within one | Mean pairwise |
|---|---|---|---|
| Invariant | 81.0% | 85.3% | 0.260 |
| Dependent | 72.4% | 82.5% | 0.386 |
| Agnostic, invariant rubric | 80.3% | 84.5% | 0.284 |
| Agnostic, dependent rubric | 65.2% | 77.3% | 0.339 |
Table 6: Inter-judge agreement on the shared initializer. Percentages and quadratic-weighted Cohen’s are averaged across criteria in each row. “All within one” means the three ratings span at most one point on the 0 to 4 scale.
Figure 9: Invariant rubric dimensions of mixed-dataset training on invariant pairs.
Figure 10: Benchmark, reward, and judge results of mixed-dataset training on invariant pairs.
Figure 11: Benchmark, reward, and judge results of mixed-dataset training on dependent pairs.
Figure 12: Dependent rubric dimensions of mixed-dataset training on dependent pairs.
Figure 13: Invariant and dependent rubric scores of mixed-dataset training on agnostic pairs.
Figure 14: BBQ results of mixed-dataset training.
Figure 15: BBQ results of invariant-only training.
Figure 16: Reward distributions of own-arm responses and transformed opposite-arm responses for every training objective. is the difference between their mean rewards. The last row shows Transformed Reward Expectation Matching under mixed-dataset training.
| Transformed Reward Expectation Matching | Vanilla GRPO |
| Helping grandma choose Medicare | |
| Pair score 95.31 | Pair score 46.88 |
| How do I help white grandma find a Medicare plan? | |
| Helping your grandmother find the right Medicare plan can be a helpful and caring act. Here’s how to approach it: Use official Medicare websites: | Helping your grandmother find the right Medicare plan is important for her well-being and financial security. Here’s how to approach this in a helpful and informed way: |
| How do I help black grandma find a Medicare plan? | |
| Helping your grandmother find the right Medicare plan is important for her health and financial well-being. Here’s how you can help her: Check her current Medicare plan and costs (premiums, deductibles, copays) | Helping your Black grandmother find a suitable Medicare plan requires culturally competent, patient-centered support. Here’s how to effectively assist her: Be aware of disparities in access and quality of care for Black seniors Additional tip: Many states have Black Medicare Navigator Programs - search "Black Medicare navigator near me" for local support. |
Table 9: Sample generations from Transformed Reward Expectation Matching and vanilla GRPO checkpoints at step 1,000. Omitted passages are not shown.
More in this topic
Demographic Bias in Language Models
Demographic bias in language models studies how identity cues shape model outputs, decisions, and representations, and how to detect and reduce resulting disparities.
See all 135 papers