ICML2025 in-context learning out-of-distribution generalization task diversity phase transition linear regression Transformer specialization

When Can In-Context Learning Generalize Out of Task Distribution?¶

会议: ICML2025
arXiv: 2506.05574
代码: GitHub
领域: LLM/NLP
关键词: in-context learning, out-of-distribution generalization, task diversity, phase transition, linear regression, transformer, specialization

一句话总结¶

通过在线性回归ICL任务上系统改变训练任务分布的覆盖范围（超球面帽的半角 \(\phi\)），发现transformer存在从"专用解"到"通用解"的sharp phase transition：当任务多样性超过临界阈值（\(\phi \gtrsim 120°\)）时，模型能泛化到整个任务空间，甚至超越贝叶斯最优估计器的OOD性能。

研究背景与动机¶

In-Context Learning 的泛化能力¶

In-context learning（ICL）是预训练transformer的一项显著能力：模型仅通过上下文中的少量示例就能推断新任务，无需重新训练。现有研究主要关注预训练任务的数量对ICL涌现的影响，但忽略了另一个维度——任务之间的相似性。

核心研究问题¶

如果模型仅在任务空间的一个子集上预训练，它能否泛化到任务空间的其余部分？

具体而言：在任务池有限且分布受限时，模型是学到一个仅在训练分布内有效的专用解（specialized solution），还是能发展出覆盖整个任务空间的通用解（general-purpose solution）？

实验范式¶

选择线性回归作为ICL的研究平台：

输入序列 \(\{x_i, y_i\}_{i=1}^n\)，其中 \(y_i = w^T x_i + \epsilon_i\)
任务由权重向量 \(w \in \mathbb{R}^d\) 定义
Transformer 的目标是在给定上下文 \(C_k = \{x_1, y_1, \dots, x_k\}\) 后预测 \(y_k\)

方法详解¶

任务分布的几何控制¶

通过超球面帽（hyperspherical cap）参数化任务分布：

\[S^{d-1}(\phi) = \{w \in S^{d-1} \mid \text{angle}(w, v) \leq \phi\}\]

\(v\) 为固定"极点"方向
\(\phi \in [0°, 180°]\) 控制覆盖范围：\(\phi = 180°\) 为整个超球面
训练分布：\(p_\phi(w) = \text{Unif}(S^{d-1}(\phi))\)
\(\phi\) 越大，训练任务越多样

测试分布¶

使用超球面带（hyperspherical band）作为测试分布：

\[B^{d-1}(\delta, \Delta\delta) = \{w \in S^{d-1} \mid \delta \leq \text{angle}(w, v) \leq \delta + \Delta\delta\}\]

\(\delta = 0°\)：测试分布在训练分布内（in-distribution）
\(\delta = 175°\)：测试分布在训练分布的"对面"（极端OOD）

训练设定¶

GPT-2 style transformer，128维hidden，10层，8头注意力
\(d = 10\) 维回归问题，每个上下文 \(n = 50\) 个样本
训练12个模型，\(\phi \in [15°, 180°]\)，步长 \(15°\)
优化目标：\(L_{\text{train}}(\theta) = \mathbb{E}_{w \sim p_\phi}\left[\frac{1}{n}\sum_{k=1}^n (T_\theta(C_k) - y_k)^2\right]\)

贝叶斯最优基准¶

给定先验 \(p_\phi(w)\)，贝叶斯最优估计器为后验均值：

\[\hat{w} = \frac{\int dw \, p(w) w \prod_{k=1}^{n-1} p(y_k | x_k, w)}{\int dw \, p(w) \prod_{k=1}^{n-1} p(y_k | x_k, w)}\]

关键性质：\(p_\phi(w)\) 在 \(S^{d-1}(\phi)\) 外无支撑，因此 \(\hat{w}\) 必然落在 \(S^{d-1}(\phi)\) 内 → 贝叶斯最优估计器无法有意义地OOD泛化。

实验关键数据¶

核心发现：专用→通用的相变¶

\(\phi\)	\(\delta = 0°\)（分布内）	\(\delta = 175°\)（极端OOD）	解类型
\(\leq 90°\)	低误差	高误差	专用解
\(\approx 120°\)	低误差	开始下降	临界点
\(\geq 120°\)	低误差	低误差	通用解

Sharp transition at \(\phi \approx 120°\)：

\(\phi < 120°\)：模型在训练分布外性能骤降
\(\phi \geq 120°\)：模型对所有测试角度 \(\delta\) 表现一致且优异
有noise时（\(\sigma^2 = 0.25\)），临界点移至 \(\phi \approx 135°\)

专用解 vs 通用解的性质¶

性质	专用解（\(\phi < 120°\)）	通用解（\(\phi \geq 120°\)）
短上下文性能	优于OLS	与OLS相当
OOD性能	差	优于贝叶斯最优
本质	拟合了强先验	学到了通用算法（类似OLS）

专用解在分布内超越OLS：通过拟合训练分布的先验，专用解在少量样本（\(k < d\)）时对分布内任务表现更好，但以牺牲OOD泛化为代价。

通用解超越贝叶斯最优：贝叶斯最优被训练先验"锁死"在 \(S^{d-1}(\phi)\) 内，而模型通过未能完美拟合贝叶斯解反而获得了OOD泛化能力。

相图：两种任务多样性的交互¶

训练480个模型（\(\phi \times N\) 的网格），其中 \(N\) 为不同training task数量。揭示三个相：

相	分布内泛化	分布外泛化	条件
In-weights Learning (IWL)	✕	✕	\(N\) 小，\(\phi\) 任意
In-distribution ICL	✓	✕	\(N\) 大，\(\phi < 120°\)
OOD ICL	✓	✓	\(N\) 大，\(\phi \geq 120°\)

相边界呈对角线结构，表明两种多样性维度之间存在权衡。

维度和深度的影响¶

因素	对临界角 \(\phi_c\) 的影响
回归维度 \(d\)（3, 5, 10, 20）	无影响，\(\phi_c \approx 120°\) 恒定
模型深度（2, 3, 10层）	无影响，\(\phi_c \approx 120°\) 恒定

临界角不随维度或深度变化，说明这不是高维几何的简单后果。

非线性回归的扩展¶

任务类型	临界角 \(\phi_c\)
线性回归	\(\approx 120°\)
逻辑回归（分类）	\(\approx 135°\)
非线性回归（单隐层ReLU网络，分别采样）	\(\approx 135°\)
非线性回归（联合采样）	\(\approx 60°\)

相变现象在非线性和分类任务中同样存在，表明这可能是ICL的普遍特征。

超球面之外的泛化¶

训练在单位超球面（\(R = 1\)）上的模型能否泛化到其他半径？

当 \(\phi \geq 45°\) 时，模型完美泛化到 \(R < 1\) 的内部球面
任务多样性不仅驱动球面上的OOD泛化，还驱动超越球面的泛化

亮点与洞察¶

相变现象的发现：首次揭示ICL泛化中存在从专用到通用的sharp phase transition，且临界点对维度和深度鲁棒
任务多样性的新维度：提出用任务空间覆盖范围（\(\phi\)）而非仅任务数量（\(N\)）来衡量多样性，两者有本质不同
超越贝叶斯最优的OOD泛化：通用解的OOD性能优于贝叶斯最优估计器，这正是因为模型未能完美拟合训练先验
三相相图：IWL → in-distribution ICL → OOD ICL 的三相结构清晰刻画了从记忆到泛化的完整图景
对LLM的启示：解释了为何LLM能在ICL中解决预训练分布外的任务——足够多样的训练数据促使模型学到通用算法而非特定先验

局限性¶

仅在简化设定中验证：线性/简单非线性回归远不能代表LLM面临的真实ICL任务复杂度
任务相似性度量有限：线性任务自然使用内积 \(w_1^T w_2\) 度量相似性，但更通用的任务缺乏对应的度量
无解析理论：所有结论均为经验性观察，缺乏解释相变机制的严格理论（如Lu et al., 2024所做的可解析模型）
模型规模较小：使用128维hidden的小型transformer，与实际LLM的规模差距巨大
未考虑concept shift：仅研究了task distribution shift，而Yadlowsky et al. (2024)已表明transformer在concept shift下通常无法泛化

评分¶

⭐⭐⭐⭐ — 实验设计优雅，发现了clean且robust的相变现象，对理解ICL泛化机制有重要理论启发。但简化设定限制了对LLM实际ICL行为的直接推广。