Walking the Tightrope: Disentangling Beneficial and Detrimental Drifts in Non-Stationary Custom-Tuning

TL;DR

The paper tackles detrimental concept drift in chain-of-thought reasoning during non-stationary reinforcement fine-tuning of multi-modal LLMs, particularly in medical imaging. It formalizes CoT as autoregressive token streams subject to non-stationary distributions and introduces Counterfactual Preference Optimization (CPO), which uses a hierarchical concept graph as an expert to generate counterfactual reasoning trajectories and disentangle beneficial domain adaptation from harmful drift. CPO optimizes a dual objective that aligns with human preferences while penalizing drift-inducing counterfactuals, enabling robust RFT in non-stationary environments. A large-scale chest radiography dataset, CXR-CounterFact (CCF) with 320{,}416 counterfactual reasoning trajectories, supports evaluation and benchmarking. Experiments demonstrate improved robustness, generalization, and coordination in radiology downstream tasks, underscoring CPO's potential for reliable medical AI under evolving data distributions.

Abstract

This paper uncovers a critical yet overlooked phenomenon in multi-modal large language models (MLLMs): detrimental concept drift within chain-of-thought (CoT) reasoning during non-stationary reinforcement fine-tuning (RFT), where reasoning token distributions evolve unpredictably, thereby introducing significant biases in final predictions. To address this, we are pioneers in establishing the theoretical bridge between concept drift theory and RFT processes by formalizing CoT's autoregressive token streams as non-stationary distributions undergoing arbitrary temporal shifts. Leveraging this framework, we propose a novel counterfact-aware RFT that systematically decouples beneficial distribution adaptation from harmful concept drift through concept graph-empowered LLM experts generating counterfactual reasoning trajectories. Our solution, Counterfactual Preference Optimization (CPO), enables stable RFT in non-stationary environments, particularly within the medical domain, through custom-tuning of counterfactual-aware preference alignment. Extensive experiments demonstrate our superior performance of robustness, generalization and coordination within RFT. Besides, we also contributed a large-scale dataset CXR-CounterFact (CCF), comprising 320,416 meticulously curated counterfactual reasoning trajectories derived from MIMIC-CXR. Our code and data are public.

Figures (5)

• Figure 1: Concept Drift in RFT's reasoning for chest diagnosis. Despite analogous occurrence probabilities of "lung opacity" (in red) and "opacity" (in blue) tokens during the CoT, non-stationarity induces significant bad distributional drift in clinical conclusions, especially the opposite diagnosis of atelectasis, cardiomegaly and pneumonia.
• Figure 2: The main contributions of our methods. (a) By formalizing autoregressive CoT generation as a stream of next-token prediction actions under the theoretical lens of concept drift, we reveal that even minor perturbations in reinforced fine-tuning can induce unpredictable distributional changes of final predicted results. (b) To disentangle detrimental drift, we introduce the concept graph that generates radiologically plausible counterfactual CoTs through controlled attribute perturbations. Green lines represent attributes that are positively correlated with the disease, while red denote they are exclusive. (c) We propose counterfactual preference optimization to drive the reinforced custom-tuning of MLLMs, enabling generalized CoT reasoning in non-stationary environments through disentanglement of beneficial domain adaptation from spurious concept drift, thereby achieving robust human-aligned decision-making via preference distillation.
• Figure 3: Structural Causal Graph.X: Inputs, Z: Prediction Results, T: Chain-of-Thought, and D: Latent Concept Drift within CoT under Non-stationary Reinforced Custom-Tuning.
• Figure 4: Non-stationarity of MIMIC-CXR with its percentage of diseases. Blue signifies patients with clinically confirmed diagnoses showing the long-tailed characteristic, while red demarcates suspected cases emphasizing the inherent uncertainty within medicine.
• Figure 5: Samples of CXR-CounterFact (CCF) Dataset.

Theorems & Definitions (1)

• Definition 2.1