Fine-Tuning Pre-trained Language Models for Robust Causal Representation Learning

2024-10-18

arXiv.org

AI Summary

To address spurious correlations in pretrained language models caused by reliance on latent confounders during single-domain fine-tuning—leading to poor out-of-distribution generalization—this paper proposes a robust representation learning framework based on causal front-door adjustment. Our method is the first to embed causal inference into the fine-tuning process, enabling causal representation disentanglement without requiring multi-domain data, distributional assumptions, or ignorability of hidden variables. Key technical components include causal decomposition of fine-tuned representations, causal-structure-aware data augmentation, and interpretable proxy variable construction. Extensive experiments on synthetic and real-world text tasks demonstrate significant improvements over state-of-the-art causal learning and domain adaptation methods. Results validate that fine-tuned representations can serve as reliable causal proxies, establishing a novel paradigm for robust generalization in single-domain settings.

Technology Category

• Machine Learning: Causal Learning
• Natural Language Processing: Safety and Robustness
• Reasoning under Uncertainty: Causality

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

The fine-tuning of pre-trained language models (PLMs) has been shown to be effective across various domains. By using domain-specific supervised data, the general-purpose representation derived from PLMs can be transformed into a domain-specific representation. However, these methods often fail to generalize to out-of-domain (OOD) data due to their reliance on non-causal representations, often described as spurious features. Existing methods either make use of adjustments with strong assumptions about lack of hidden common causes, or mitigate the effect of spurious features using multi-domain data. In this work, we investigate how fine-tuned pre-trained language models aid generalizability from single-domain scenarios under mild assumptions, targeting more general and practical real-world scenarios. We show that a robust representation can be derived through a so-called causal front-door adjustment, based on a decomposition assumption, using fine-tuned representations as a source of data augmentation. Comprehensive experiments in both synthetic and real-world settings demonstrate the superior generalizability of the proposed method compared to existing approaches. Our work thus sheds light on the domain generalization problem by introducing links between fine-tuning and causal mechanisms into representation learning.

Problem

• Adapting to latent-confounded shifts in AI models
• Identifying causal features in confounded fine-tuning data
• Improving model robustness against spurious correlations

Innovation

• Causal fine-tuning addresses latent-confounded shifts
• Decomposes inputs into spurious and causal features
• Uses pre-trained models for automatic adaptation