CausalMix: Data Mixture as Causal Inference for Language Model Training¶

会议: NeurIPS 2026
arXiv: 2607.01104
OpenReview: https://openreview.net/forum?id=B9ShLVNHQD
论文: https://arxiv.org/abs/2607.01104
领域: 因果推理
关键词: 因果推断, 数据配比, 监督微调, 双重机器学习, 条件平均处理效应

一句话总结¶

将大语言模型监督微调（SFT）阶段的数据配比优化重构为状态条件下的因果边际收益估计问题，利用双重机器学习（DML）与因果森林剥离数据池内在质量的混杂偏差，实现了无需重新运行代理实验即可跨数据规模、跨模型架构外推的最优数据配比。

研究背景与动机¶

在大语言模型（LLM）的监督微调（SFT）阶段，指令遵循、数学推理、代码编写、知识问答以及安全对齐等多领域数据的配比比例对模型的最终综合能力有着决定性影响。然而，在大规模连续单纯形空间中穷举搜索最优配比极其昂贵。现有的自动化数据配比方法（如 RegMix、DoReMi、DMO 等）大多基于预训练视角的验证损失最小化或静态代理模型拟合，它们建立在一个隐式假设之上：候选数据池的分布与质量是静态恒定的。当实际生产中的数据池发生漂移、难度变化或规模扩展时，这些依赖全局经验映射的方法往往失效，被迫耗费高昂算力从头训练大量代理模型。

这种失效的核心矛盾在于：数据池本身的内在特征（如整体复杂度、固有难度、样本风格质量）与下游评测得分之间存在强烈的基线关联。直接用监督回归拟合「配比比例到绝对得分」的映射，本质上混淆了「候选数据本身的质量水平」与「调整领域配比所带来的真实因果收益」。标准回归模型受制于正则化偏差与共线性，学到的往往是特定数据池下的虚假相关性，既无法解耦领域之间的因果迁移效应，更无法跨越到全新的未见数据池中。

本文的切入角度是打破将配比优化视为黑盒超参数搜索的传统视角，将其重构为条件平均处理效应（CATE）的因果推断问题。通过将数据池的统计特征形式化为协变量、将领域配比视为连续干预，并利用正交因果学习剔除数据状态的混杂影响。核心 idea：将 SFT 数据配比转化为状态条件下的因果边际收益估计，利用双重机器学习与因果森林正交剥离数据池质量混杂偏差，并通过 KKT 条件解析求解单纯形最优配比，实现跨规模与跨模型的高泛化数据配比策略。

方法详解¶

整体框架¶

CausalMix 的核心流程由历史代理实验收集、状态协变量提取、对数配比参数化、双重机器学习正交因果估计，以及最优单纯形策略提取五个核心环节构成。系统首先在小参数规模的代理模型（如 Qwen2.5-0.5B）上运行历史微调实验，收集「数据池状态特征-领域配比-下游评测表现」三元组；随后利用双重机器学习框架剥离数据状态对干预变量和结果变量的混杂影响，训练因果森林模型拟合状态条件下的边际收益；最后根据目标数据池的特征，通过闭式解析解或局部候选搜索，将边际收益映射为满足概率单纯形约束的最终领域混合权重。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：历史代理实验三元组 (X, T, Y)<br/>与目标数据池统计特征 X_tar"] --> B["对数配比表征与状态条件潜在结果建模<br/>单纯形几何转换 Z=log(T+ε) 与部分线性逼近"]
    B --> C["基于双重机器学习的正交因果效应估计<br/>残差化剥离协变量偏差，因果森林求解 CATE"]
    C --> D["单纯形闭式解析策略与局部 Bagging 搜索<br/>KKT 闭式解 (CausalMix-A) 或 Top 候选集成 (CausalMix-S)"]
    D --> E["输出：最优领域配比权重 T*<br/>指导大规模模型全量 SFT 训练"]

关键设计¶

1. 对数配比表征与状态条件潜在结果建模：解决单纯形非线性边际递减与基线混杂问题 在数据配比任务中，直接以原始概率单纯形向量 $T = (T_1, \dots, T_K)$（满足 $T_k \ge 0, \sum_k T_k = 1$）作为干预变量，难以直接刻画神经缩放定律所揭示的「边际收益递减」效应，同时单纯形的求和归一约束打破了变量独立性。为此，方法引入成分数据分析中的几何变换，采用对数配比参数化： $$Z = \log(T + \varepsilon)$$ 其中 $\varepsilon > 0$ 为极小平滑常数。对于观测到的状态协变量 $X$（涵盖数据池的复杂度、难度与质量），定义潜在结果条件响应曲面 $\mu(x, z) = \mathbb{E}[Y(t) \mid X=x]$。由于历史代理实验样本有限，无法在高维连续空间拟合完整曲面，方法采用局部部分线性逼近： $$\mu(x, Z) \approx g(x) + \theta_0(x)^\top Z$$ 其中 $g(x)$ 刻画由数据池自身固有质量决定的基线表现，而 $\theta_0(x) \in \mathbb{R}^K$ 则代表目标数据状态 $x$ 下各领域的因果边际收益（广义多维连续 CATE）。若 $\theta_{0,k}(x) > 0$，表明在当前数据状态下增加该领域比例能产生正向迁移；若为负值则提示存在负迁移冲突。

2. 基于双重机器学习的正交因果效应估计：剥离数据池内在质量对边际收益的混杂偏差 直接用标准监督回归拟合绝对得分 $Y$ 会将基线能力 $g(X)$ 与配比处理效应混杂在一起，导致严重的正则化偏差。CausalMix 引入双重机器学习（Double Machine Learning, DML）框架，通过构造正交化残差来隔离纯粹的处理效应。首先定义两个讨厌函数（nuisance functions）： $$m_0(X) = \mathbb{E}[Y \mid X], \quad e_0(X) = \mathbb{E}[Z \mid X]$$ 并通过多折交叉拟合（cross-fitting）分别预测结果期望与对数配比期望，从而构建正交残差： $$\widetilde{Y} = Y - m_0(X), \quad \widetilde{Z} = Z - e_0(X)$$ 此时边缘收益满足正交关系 $\widetilde{Y} \approx \theta_0(X)^\top \widetilde{Z}$。模型通过最小化 Neyman 正交损失（R-loss）训练因果效应模型： $$\hat{\theta} = \arg\min_\theta \sum_i \left( \widetilde{Y}_i - \theta(X_i)^\top \widetilde{Z}_i \right)^2$$ 在第一阶段预测器上，采用双 LightGBM 结构分别拟合 $m_0(X)$ 与 $e_0(X)$；在因果效应估计器上，选用因果森林（CausalForestDML）。因果森林利用非参数树形自适应递归划分，能够精准捕捉连续多维协变量空间中的异质性效应与饱和阈值，在 R-Scorer 评估中显著优于线性 DML 模型。

3. 单纯形闭式解析策略与局部 Bagging 搜索：将连续边际收益严谨转化为单纯形配比 估计出目标数据池状态下的边际收益向量 $\hat{\theta}(X_{\text{tar}})$ 后，核心挑战在于如何将其合法映射回满足非负与归一化约束的概率单纯形空间。基于 Level-Log 收益最大化目标 $\max_T \sum_k \hat{\theta}_k \log(T_k)$ 受限于 $\sum_k T_k = 1, T_k \ge 0$，论文推导了基于 Karush–Kuhn–Tucker（KKT）条件的最优解析解（CausalMix-A）： $$T_k^{\mathrm{A}} = \frac{[\hat{\theta}_k(X_{\mathrm{tar}})]_+}{\sum_{j=1}^K [\hat{\theta}_j(X_{\mathrm{tar}})]_+}, \quad [a]_+ = \max(a, 0)$$ 数学证明表明，对所有边际收益非正的领域（$\hat{\theta}_k \le 0$），最优解严格取边界值 0；对正收益领域则按边际贡献比例严格归一化。此外，为缓解单点预测在未知分布下的估计方差，论文还提出了基于采样的搜索变体（CausalMix-S）：从狄利克雷分布中抽取 100,000 个候选配比，输入拟合好的因果模型计算预估得分，并对得分最高的 Top-100 候选配比在原始单纯形空间取平均： $$T^{\mathrm{S}} = \frac{1}{K_{\mathrm{top}}} \sum_{m \in \mathrm{Top}} T^{(m)}$$ 这种局部 Bagging 集成有效平滑了噪声，提供了更稳健的域外泛化配比。

损失函数 / 训练策略¶

在数据协变量的选取上，研究基于 OpenDataArena 的 30 维特征进行了系统消融，最终收敛至三维核心协变量体系： 1. 复杂度指标 HES（High Entropy Sum）：统计 Qwen3-8B 推理链中熵值最高的前 0.5% token 的熵值和，精准锚定复杂推理决策节点； 2. 难度指标 Normalized_Loss：由 Qwen3-8B 计算样本归一化交叉熵损失，表征样本的可预测性与训练难度； 3. 质量指标 Writing_Style：基于 QuRater-1.3B 评估文本的清晰度、连贯性与表达规范性。在代理模型训练阶段，使用 Qwen2.5-0.5B 在 512 个采样的 100K 子数据集上进行全参数微调（学习率 2.0e-5，余弦退火，序列长度 4096，共 3 epochs），总计算量约 $5.53 \times 10^{20}$ FLOPs。完成因果模型拟合后，后续针对 800K 数据池及 7B 模型乃至全新的 LongCoT 任务，均直接外推推断配比，无需增加任何额外代理训练成本。

实验关键数据¶

主实验¶

主实验在 tulu-3-sft-mixture 数据集上展开，覆盖通用指令遵循（IF）、数学推理（Math）、代码编程（Coding）、知识召回（Knowledge）和安全合规（Safety）五大领域。评测涵盖开发集（$\mathcal{S}_{\mathrm{Dev}}$，10 个评测基准）和未见测试集（$\mathcal{S}_{\mathrm{Uns}}$，6 个 OOD 评测基准）。实验在 Qwen2.5-0.5B（不同数据规模）及 Qwen2.5-7B 上全面验证。

模型与数据规模	方法	Knowledge	Reasoning	Math	Coding	IF	Safety	$\mathbf{Avg}_{\mathrm{Dev}}$	$\mathbf{Avg}_{\mathrm{Uns}}$
Qwen2.5-0.5B, 100K	Equal	27.65	29.80	23.32	33.00	33.83	14.57	27.03	24.43
Qwen2.5-0.5B, 100K	Grid	28.46	31.98	30.41	33.50	20.15	23.38	27.98	23.45
Qwen2.5-0.5B, 100K	RegMix	27.29	30.04	25.53	31.28	17.01	32.56	27.28	26.24
Qwen2.5-0.5B, 100K	DoReMi	29.81	31.70	22.90	33.40	35.30	25.34	29.74	24.12
Qwen2.5-0.5B, 100K	DMO	28.86	30.91	22.51	30.68	35.86	26.19	29.17	26.08
Qwen2.5-0.5B, 100K	CausalMix-A	28.37	31.31	24.67	29.57	38.63	26.93	29.91	23.42
Qwen2.5-0.5B, 100K	CausalMix-S	27.27	30.23	27.77	34.94	34.94	14.81	27.85	25.90
Qwen2.5-0.5B, 800K	Equal	28.07	29.28	21.59	36.06	46.95	25.09	31.78	25.64
Qwen2.5-0.5B, 800K	RegMix	24.90	29.94	30.25	25.46	25.32	22.52	26.40	22.82
Qwen2.5-0.5B, 800K	DMO	28.07	31.88	22.42	26.90	41.59	41.37	32.04	26.49
Qwen2.5-0.5B, 800K	CausalMix-A	27.93	29.68	23.56	27.76	43.81	50.92	33.94	25.02
Qwen2.5-0.5B, 800K	CausalMix-S	28.31	30.96	27.64	30.97	42.51	36.47	32.81	25.04
Qwen2.5-7B, 800K	Equal	60.85	64.55	59.03	53.61	68.58	53.49	60.02	49.55
Qwen2.5-7B, 800K	RegMix	59.60	68.12	63.37	55.76	58.04	55.94	60.14	48.12
Qwen2.5-7B, 800K	DoReMi	58.02	63.20	57.35	57.28	68.21	52.75	59.47	46.19
Qwen2.5-7B, 800K	DMO	59.15	63.70	60.62	54.05	70.24	54.35	60.35	48.98
Qwen2.5-7B, 800K	CausalMix-A	57.14	64.03	58.51	65.52	68.21	57.65	61.84	49.09
Qwen2.5-7B, 800K	CausalMix-S	59.35	62.88	58.63	64.43	67.47	60.95	62.28	47.98

此外，在 LongCoT 长推理跨架构迁移实验中（Qwen3-4B-Base 模型，20K 数据，AM-Thinking-v1-Distilled 数据集），CausalMix 将历史学习到的因果策略直接应用于未见模型与长链数据，结果如下：

方法	GSM8K	MATH	$\mathbf{Avg}_{\mathrm{Math}}$	HumanEval	MBPP	$\mathbf{Avg}_{\mathrm{Code}}$	$\mathbf{Avg}_{\mathrm{Total}}$
Equal	90.45	56.78	73.62	59.76	48.20	53.98	63.80
Grid	87.34	61.20	74.27	62.80	47.60	55.20	64.74
RegMix	89.61	40.80	65.21	61.59	53.60	57.60	61.40
DoReMi	88.55	42.22	65.39	63.41	53.80	58.61	62.00
ODM	88.32	41.16	64.74	63.41	42.20	52.81	58.77
DMO	89.61	54.38	72.00	54.88	55.00	54.94	63.47
CausalMix	88.86	60.58	74.72	62.20	55.00	58.60	66.66

消融实验¶

消融实验重点验证了因果协变量 $X$ 的引入与 DML 正交化步骤（Orthogonalization）的必要性。对比变体包括：(1) w/o X：去除协变量，直接学习静态配比映射 $\hat{Y} = g(T)$；(2) w/o Orth.：绕过正交化，直接拼接协变量与配比进行监督回归 $\hat{Y} = f(X, T)$。

模型与规模	实验配置	Knowledge	Reasoning	Math	Coding	IF	Safety	$\mathbf{Avg}$	说明
Qwen2.5-0.5B, 800K	w/o X	29.27	31.39	29.97	33.41	39.37	36.35	33.29	丢失数据状态感知
Qwen2.5-0.5B, 800K	w/o Orth.	27.41	31.29	24.74	31.90	41.04	37.82	32.66	存在严重正则化偏差
Qwen2.5-0.5B, 800K	CausalMix-A	27.93	29.68	23.56	27.76	43.81	50.92	33.94	闭式最优解析解
Qwen2.5-0.5B, 800K	CausalMix-S	28.31	30.96	27.64	30.97	42.51	36.47	32.81	Top-100 集成搜索
Qwen2.5-7B, 800K	w/o X	60.45	63.95	61.16	55.62	69.69	56.92	61.30	跨规模迁移能力退化
Qwen2.5-7B, 800K	w/o Orth.	59.50	64.66	60.45	45.82	68.76	58.75	59.65	混淆基线导致大幅掉点
Qwen2.5-7B, 800K	CausalMix-A	57.14	64.03	58.51	65.52	68.21	57.65	61.84	代码能力提升显著 (+19.7)
Qwen2.5-7B, 800K	CausalMix-S	59.35	62.88	58.63	64.43	67.47	60.95	62.28	综合得分最高

在因果模型与一阶段预测器的模型选择上，基于 R-Scorer 评估： - 因果效应模型选择：CausalForestDML 取得最高的 RScore (+0.1683)，明显超越 LinearDML (+0.1445)、DML_Poly3_Lasso (+0.1533) 与 SparseLinearDML (-1.7065)。 - 第一阶段模型选择：使用双 LightGBM 结构预测 $Y$ 与 $T$ 时表现最佳（RScore 0.1683，耗时 12.9s），而纯 RandomForest 结构 RScore 降为 -0.2840。

关键发现¶

直接回归的正则化偏差甚至劣于纯配比拟合：在 7B 模型上，w/o Orth. 的综合平均分仅 59.65，甚至显著低于完全不使用协变量的 w/o X（61.30）。这证明简单将数据特征与配比拼接做常规监督学习时，标准损失惩罚会迫使模型过度依赖特征本身而压制微弱的配比因果信号，正交化是因果效应识别的绝对先决条件。
解析解与搜索解的互补性：CausalMix-A 在开发集上往往能推导出极致的极端优化解（如在 0.5B 800K 上取得 33.94 的最高开发集均分），但在未见分布（$\mathcal{S}_{\mathrm{Uns}}$）上，CausalMix-S 通过 Top-100 候选配比的局部平均，有效平滑了个别维度的估计方差，具备更优异的泛化鲁棒性。
异质性效应的树解释（CATE Interpreter）揭示技能冲突：树解释器分析表明，通用指令遵循（IF）在所有数据子空间中均呈现稳定的正边际收益，是对齐的基础支撑；但在高 Normalized_Loss 与高 HES（高难度、高复杂度）的数据区域，Knowledge 数据却呈现明显的负边际效应，实证验证了事实记忆与深层逻辑推理之间的技能冲突（skill conflict）。此外，在低质量数据区域（低 Writing_Style 与低 HES），盲目增加数学和代码数据会引入过多噪声导致负迁移；而在中高质量区域，数学、代码与安全数据会产生极强的正向协同增益。

亮点与洞察¶

将数据配比升维至状态条件因果推断：彻底打破了以往将数据混合视作纯粹经验拟合或静态损失优化的局限，利用 DML 框架在理论上严格消除了数据池自身质量这一主要混杂变量，使模型学到的是可迁移的配比变化边际收益率，而非死记特定数据集的静态比例。
KKT 条件下的单纯形闭式解析解证明：严谨证明了在 Level-Log 效应模型下，边际收益非正的领域最优权重必为 0，而正收益领域权重等价于 ReLU 激活后的一范数归一化，为启发式调参提供了坚实的数学闭式依据。
可复用的 CATE 解释器工具链：利用因果树解释器对微调多任务数据交互进行可视化，不仅能指导数据配比，还能在数据合成与清洗阶段为判定何种质量阈值下应注入代码/数学数据提供量化依据。

局限与展望¶

代理实验的离散采样规模依赖：当前因果模型拟合仍依赖 512 次 0.5B 代理微调实验作为基础元数据集（Meta-dataset）。虽然一次拟合后可终身多次外推，但初次启动时的冷启动实验成本（约 $5.53 \times 10^{20}$ FLOPs）对于计算资源极其受限的团队依然是一道门槛。
协变量维度受限于元数据规模：由于样本量为 512，协变量数量过多时会受到维度灾难影响导致 RScore 下降（因而最终仅保留了 3 个核心指标）。未来若能构建跨机构的大规模历史训练轨迹数据库，将能支持更多细粒度特征的因果建模。
静态配比到动态退火调度（Curriculum）的拓展：当前模型输出的是全周期静态配比权重，尚未考虑不同训练阶段（如前期重基础语言分布、后期重长链推理）的动态退火混合策略，未来可将因果干预扩展为时变动态策略。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将因果推断与双重机器学习（DML）严格引入 LLM SFT 数据配比优化，理论建模与 KKT 解析解推导非常优雅。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖 0.5B 到 7B、100K 到 800K 数据规模，且成功 zero-shot 迁移至 Qwen3-4B 长链推理数据，消融与因果模型选择细致入微。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备，从因果识别条件、正交损失到策略提取逻辑链条极其紧凑连贯。
价值: ⭐⭐⭐⭐⭐ 解决了微调阶段数据池漂移导致重复调参的工业界核心痛点，兼具卓越的跨模型外推实用性与深刻的技能冲突可解释性。