Mitigating Spurious Correlations in LLMs via Causality-Aware Post-Training

TL;DR

This paper introduces causality-aware post-training (CAPT) to reduce spurious correlations in large language models, improving their out-of-distribution generalization without extensive fine-tuning.

Contribution

It proposes a novel causality-aware post-training method that decomposes biased predictions into unbiased steps, enhancing model robustness and generalization.

Findings

CAPT improves OOD performance of 3B-scale LLMs.

CAPT outperforms traditional fine-tuning with only 100 samples.

CAPT enhances generalization without additional fine-tuning biases.

Abstract

While large language models (LLMs) have demonstrated remarkable capabilities in language modeling, recent studies reveal that they often fail on out-of-distribution (OOD) samples due to spurious correlations acquired during pre-training. Here, we aim to mitigate such spurious correlations through causality-aware post-training (CAPT). By decomposing a biased prediction into two unbiased steps, known as \textit{event estimation} and \textit{event intervention}, we reduce LLMs' pre-training biases without incurring additional fine-tuning biases, thus enhancing the model's generalization ability. Experiments on the formal causal inference benchmark CLadder and the logical reasoning dataset PrOntoQA show that 3B-scale language models fine-tuned with CAPT can outperform both traditional SFT and larger LLMs on in-distribution (ID) and OOD tasks using only 100 ID fine-tuning samples, demonstrating the effectiveness and sample efficiency of CAPT.