上下文线性回归

引言

Transformer架构在自然语言处理和计算机视觉领域取得了巨大成功，其核心能力之一是上下文学习（In-Context Learning, ICL）——无需参数更新，仅通过输入序列中的示例即可学习新任务¹。近年来，研究表明经过训练的Transformer能够隐式地发现并执行经典数值算法，这一发现引发了学界的广泛关注²。

Transformer作为通用学习者的发现

Garg等人（2022）和Akyürek等人（2023）首先观察到，经过训练的Transformer可以在上下文中学习简单的函数类³⁴。von Oswald等人（2023）进一步证明，对于线性回归任务，线性Transformer的每一层都在执行类似梯度下降的迭代算法，他们称之为GD++⁵。

最近，Lutz等人（2025）在NeurIPS 2025上发表的工作「Linear Transformers Implicitly Discover Unified Numerical Algorithms」中取得了突破性进展：他们训练一个线性注意力Transformer完成矩阵补全任务，通过代数展开（algebraic unrolling）揭示出同一套参数无关的迭代更新规则（EAGLE算法）在三种不同的计算模式下（集中式、分布式、计算受限）均能出现²。

上下文学习的数学框架

在上下文学习的框架下，输入序列由数据点对 $(x_i,y_i)$ 组成，其中 $x_i\in {\mathbb{R}}^d$ 是特征向量，$y_i\in \mathbb{R}$ 是对应的标签。模型需要根据这些示例预测查询点 $x_t$ 对应的标签 $y_t$。

# 上下文学习的典型输入格式
# 每个token包含 (x_i, y_i) 对
tokens = [
    (x_1, y_1),  # 示例1
    (x_2, y_2),  # 示例2
    ...
    (x_n, y_n),  # 示例n
    (x_t, 0)     # 查询token，y初始化为0
]

线性Transformer的前向传播可表示为：

$$\Delta e_i=W_P\left(\sum _{j=1}^n⟨W_Q{e}_i,W_K{e}_j⟩W_V{e}_j\right)$$

其中 $W_Q,W_K,W_V,W_P$ 是注意力层的参数矩阵。通过合并参数，可以得到更简洁的形式：

$$\Delta e_i=\sum _{j=1}^n(e_j^{\top }Q{e}_i)P{e}_j$$

线性回归的隐式求解

正规方程

给定 $n$ 个训练样本 $(x_i,y_i)$，线性回归的目标是找到权重向量 $w^{\ast }$ 最小化均方误差：

$$L(w)=\sum _{i=1}^n(y_i-⟨w,x_i⟩{)}^2$$

正规方程给出闭式解：

$$w^{\ast }={\Sigma }^{-1}\alpha$$

其中：

• $\Sigma ={\sum }_{i=1}^n{x}_i{x}_i^{\top }$ 是特征矩阵的协方差
• $\alpha ={\sum }_{i=1}^n{y}_i{x}_i$ 是标签-特征的加权和

// 正规方程求解线性回归
vector<double> solve_normal_equation(const vector<vector<double>>& X, 
                                      const vector<double>& y) {
    int n = X.size(), d = X[0].size();
    // 计算 Sigma = X^T X
    vector<vector<double>> Sigma(d, vector<double>(d, 0));
    for (int i = 0; i < n; ++i)
        for (int k = 0; k < d; ++k)
            for (int j = 0; j < d; ++j)
                Sigma[k][j] += X[i][k] * X[i][j];
    
    // 计算 alpha = X^T y
    vector<double> alpha(d, 0);
    for (int i = 0; i < n; ++i)
        for (int k = 0; k < d; ++k)
            alpha[k] += y[i] * X[i][k];
    
    // w = Sigma^{-1} * alpha (需要矩阵求逆)
    return matrix_inverse(Sigma) * alpha;
}

梯度下降法

当数据规模较大时，直接求逆的计算成本为 $O(d^3)$。梯度下降法提供了一种迭代求解方案：

$$w_{t+1}=w_t-\eta \nabla L(w_t)=w_t-\eta \sum _{i=1}^n(y_i-⟨w_t,x_i⟩)x_i$$

梯度下降的收敛速度依赖于条件数 $\kappa (\Sigma )$，通常需要 $O(\kappa \log (1/ϵ))$ 次迭代才能达到精度 $ϵ$。

def gradient_descent(X, y, lr=0.01, max_iter=1000):
    n, d = X.shape
    w = np.zeros(d)
    
    for _ in range(max_iter):
        # 计算预测误差
        residual = y - X @ w
        # 计算梯度
        gradient = X.T @ residual
        # 更新权重
        w = w + lr * gradient
        
    return w

共轭梯度法

共轭梯度法（Conjugate Gradient, CG）是求解线性方程组和最小二乘问题的一阶方法，其迭代复杂度为 $O(\sqrt{\kappa }\log (1/ϵ))$，在条件数较大时优于普通梯度下降⁶。

def conjugate_gradient(A, b, max_iter=None):
    """求解 Ax = b，其中A是对称正定矩阵"""
    x = np.zeros_like(b)
    r = b - A @ x
    p = r.copy()
    rsold = r @ r
    
    for i in range(max_iter or len(b)):
        Ap = A @ p
        alpha = rsold / (p @ Ap)
        x = x + alpha * p
        r = r - alpha * Ap
        rsnew = r @ r
        
        if np.sqrt(rsnew) < 1e-8:
            break
            
        beta = rsnew / rsold
        p = r + beta * p
        rsold = rsnew
        
    return x

Transformer的隐式执行

权重向量的维护

Vladymyrov等人（2024）证明，线性Transformer的每一层都维护着一个隐式线性回归问题的权重向量⁷。具体来说，对于第 $l$ 层的输出 $(x_i^l,y_i^l)$，存在参数 $M^l,u^l,w^l,a^l$ 使得：

$$x_i^{l+1}=M^l{x}_i+y_i{u}^l$$ $$y_i^{l+1}=a^l{y}_i-⟨w^l,x_i⟩$$

对于查询token $x_t$：

$$y_t^{l+1}=-⟨w^l,x_t⟩$$

这意味着每层都在执行一种预处理梯度下降（Preconditioned Gradient Descent）的变体。

预处理机制

GD++算法（von Oswald等人，2023）展示了一个关键的预处理机制⁵：

$$x_i^{l+1}=(I+{\omega }_{xx}^l{\Sigma }^l)x_i^l$$ $$y_i^{l+1}=y_i^l+{\omega }_{yx}^l⟨{\alpha }^l,x_i^l⟩$$

其中 ${\Sigma }^l={\sum }_{i=1}^n{x}_i^l(x_i^l{)}^{\top }$，${\alpha }^l={\sum }_{i=1}^n{y}_i^l{x}_i^l$。

关键发现：GD++实际上实现了二阶收敛！Vladymyrov等人的定理5.1表明，对于最小二乘问题，GD++可以在 $O(\log \kappa +\log \log (1/ϵ))$ 层内达到精度 $ϵ$⁷。

# GD++ 算法的预处理机制示意
class GDPlusPlus:
    def __init__(self):
        self.w = None  # 隐式权重向量
        
    def forward(self, X, y, X_query):
        """X: (n, d), y: (n,), X_query: (d,)"""
        n, d = X.shape
        w = np.zeros(d)  # 初始化权重
        
        for layer in range(self.num_layers):
            # 计算统计量
            Sigma = X.T @ X  # d x d
            alpha = X.T @ y  # d
            
            # 预处理步骤：更新x（特征）
            # x' = (I + ω_xx * Sigma) @ x
            
            # 梯度步骤：更新y（残差）
            residual = y - X @ w
            gradient = X.T @ residual
            w = w + self.lr * gradient
            
        return X_query @ w

与在线学习的联系

Transformer的逐层迭代可以类比为在线学习中的优化过程：每一层接收前一层的信息（对应”历史梯度”），然后进行一步更新。这种类比揭示了Transformer如何通过堆叠注意力层来”模拟”迭代优化算法。

Lutz等人（2025）的EAGLE算法进一步深化了这一理解²。通过训练Transformer完成矩阵补全任务，他们发现权重重构后揭示出统一的迭代规则，该规则能够：

• 在集中式设置中实现二阶收敛
• 在分布式设置中减少通信复杂度
• 在计算受限设置中保持准确性

预处理共轭梯度法

PCG算法回顾

预处理共轭梯度法（Preconditioned Conjugate Gradient）是求解对称正定线性系统 $Ax=b$ 的经典方法。其核心思想是通过一个预处理矩阵 $M$ 来改善矩阵 $A$ 的条件数：

$$M^{-1}Ax=M^{-1}b$$

PCG的基本迭代为：

$${\alpha }_k=\frac{r_k^{\top }{r}_k}{p_k^{\top }A{p}_k}$$ $$x_{k+1}=x_k+{\alpha }_k{p}_k$$ $${\beta }_k=\frac{r_{k+1}^{\top }{r}_{k+1}}{r_k^{\top }{r}_k}$$ $$p_{k+1}=r_{k+1}+{\beta }_k{p}_k$$

def preconditioned_cg(A, b, M=None, max_iter=1000, tol=1e-8):
    """
    预处理共轭梯度法
    A: 对称正定矩阵
    b: 右端向量
    M: 预处理矩阵（通常为A的近似逆）
    """
    n = len(b)
    x = np.zeros(n)
    r = b - A @ x
    z = r if M is None else np.linalg.solve(M, r)
    p = z.copy()
    
    for k in range(max_iter):
        Ap = A @ p
        pAp = p @ Ap
        
        if abs(pAp) < 1e-12:  # 避免除零
            break
            
        alpha = (r @ z) / pAp
        x = x + alpha * p
        r = r - alpha * Ap
        z_new = r if M is None else np.linalg.solve(M, r)
        
        if np.linalg.norm(r) < tol:
            break
            
        beta = (z_new @ r) / (z @ r)
        p = z_new + beta * p
        z = z_new
        
    return x

隐式PCG的数学推导

EAGLE算法（Lutz等人，2025）可以被理解为一种隐式的预处理共轭梯度变体²。考虑矩阵补全问题：

$$X=\left[\begin{array}{cc}A& C\\ B& D\end{array}\right],\text{rank}(X)=\text{rank}(A)$$

目标是根据可见块 $A,B,C$ 预测缺失块 $D$。Nyström近似给出：

$${\hat{D}}^{\ast }=BA(A{A}^{\top }{)}^{†}C$$

EAGLE的核心更新规则为：

$$\tilde{A}\leftarrow AS,\tilde{B}\leftarrow BS$$ $$\rho \leftarrow \Vert \tilde{A}{\Vert }_2^{-2}$$ $$A^{\prime }\leftarrow A-\eta \rho \tilde{A}{\tilde{A}}^{\top }\tilde{A}{S}^{\top }$$ $$B^{\prime }\leftarrow B-\eta \rho \tilde{B}{\tilde{A}}^{\top }\tilde{A}{S}^{\top }$$ $$C^{\prime }\leftarrow C-\gamma \rho \tilde{A}{\tilde{A}}^{\top }C$$ $$D^{\prime }\leftarrow D+\gamma \rho \tilde{B}{\tilde{A}}^{\top }C$$

其中 $S$ 是正交草图矩阵，$\eta \approx 1$，$\gamma \approx 1.9$。

def EAGLE_update(A, B, C, D, S, eta=1.0, gamma=1.9):
    """
    EAGLE算法的核心更新
    A, B, C, D: 矩阵块
    S: 正交草图矩阵
    """
    # 计算草图
    A_tilde = A @ S
    B_tilde = B @ S
    
    # 计算缩放因子
    rho = 1.0 / (np.linalg.norm(A_tilde, 2) ** 2)
    
    # 更新各块
    A_new = A - eta * rho * A_tilde @ A_tilde.T @ A @ S.T
    B_new = B - eta * rho * B_tilde @ A_tilde.T @ A @ S.T
    C_new = C - gamma * rho * A_tilde @ A_tilde.T @ C
    D_new = D + gamma * rho * B_tilde @ A_tilde.T @ C
    
    return A_new, B_new, C_new, D_new

收敛性分析

定理1（集中式设置的二阶收敛）²：

对于任意矩阵 $X$，设 ${\hat{D}}^{\ast }$ 是 $D$ 的Nyström估计，$\kappa =\kappa (A_0)$。若 $\eta =1/3$，$\gamma =1$，则存在

$$L=O\left(\log \kappa +\log \log \left(\frac{1}{ϵ\sqrt{d^{\prime }}\Vert W^{\ast }{\Vert }_F\Vert C{\Vert }_F}\right)\right)$$

使得对所有 $l\ge L$：

$$\Vert D_l-{\hat{D}}^{\ast }{\Vert }_F\le ϵ$$

关键洞察：EAGLE的迭代与经典的Newton-Schulz矩阵求逆方法密切相关。考虑归一化后的迭代：

$${\bar{A}}_{l+1}=\frac{1}{2}(3I-{\bar{A}}_l{\bar{A}}_l^{\top }){\bar{A}}_l$$

这正是计算矩阵符号函数的Newton-Schulz迭代²。

秩限制注意力

低秩近似的意义

在计算受限的场景中，注意力矩阵的秩被限制为 $r\ll n$。这种约束带来两个优势：

1. 计算效率：每层复杂度从 $O(n^{2}d)$ 降低到 $O(nrd)$
2. 通信效率：分布式设置中的通信量从 $O(n^2)$ 降低到 $O(n)$

Lutz等人发现，即使在这种低秩约束下，Transformer仍然能够学习到有效的算法²。

近似精度分析

分布式设置中的多样性指数²：

定义数据多样性指数 $\alpha$ 为：

$$\alpha ={\lambda }_{\min }\left(\frac{1}{M}\sum _{\mu =1}^M{P}_{\mu }\right)$$

其中 $P_{\mu }$ 是 $A_{\mu }$ 值域上的正交投影。

定理2（分布式收敛）：

设 ${\kappa }_{\max }={\max }_{\mu }\kappa (A_{\mu })$。在无噪声情况下，存在

$$L=O\left(\log \left({\kappa }_{\max }+{\alpha }^{-1}\log \left(\frac{\sqrt{d^{\prime }}\Vert C{\Vert }_F\Vert W^{\ast }{\Vert }_F}{ϵ}\right)\right)\right)$$

使得 $\Vert D_l-D^{\ast }{\Vert }_F\le ϵ$。

def compute_diversity_index(A_list):
    """
    计算分布式设置中的数据多样性指数
    A_list: 各机器上的数据矩阵列表
    """
    M = len(A_list)
    d = A_list[0].shape[0]
    
    # 计算各机器的投影矩阵
    P_sum = np.zeros((d, d))
    for A_mu in A_list:
        Q, _ = np.linalg.qr(A_mu)
        P_sum += Q @ Q.T
    
    # 计算最小特征值
    P_avg = P_sum / M
    eigenvalues = np.linalg.eigvalsh(P_avg)
    alpha = eigenvalues[0]  # 最小特征值
    
    return alpha

计算受限设置中的草图分析：

在秩限制注意力中，草图矩阵 $S\in {\mathbb{R}}^{n\times r}$ 呈现出随机正交矩阵的性质。其特征值分布和系数分布与 $S{S}^{\top }$ 的统计特性高度吻合，这解释了为什么低秩约束不会显著损害算法性能。

实验验证

矩阵补全任务

Lutz等人（2025）的实验设置²：

• 数据生成：构造秩为 $s$ 的矩阵 $X=R_1{R}_2^{\top }/\sqrt{s}$，其中 $R_1,R_2$ 的行从 $N(0,\Sigma )$ 采样
• 参数：$n=d=18$，$n^{\prime }=d^{\prime }=2$，$s=10$，$\alpha =0.7$
• 噪声：添加方差为 $0.01$ 的高斯噪声

实验结果显示，4层Transformer在三种设置下均实现了 $10^{-50}$ 量级的MSE降低。

# 矩阵补全任务的典型实验配置
class MatrixCompletionExperiment:
    def __init__(self):
        self.n = 18      # 训练样本数
        self.d = 18      # 特征维度
        self.n_prime = 2 # 输出维度
        self.rank = 10   # 矩阵秩
        self.alpha = 0.7 # 条件数参数
        
    def generate_low_rank_matrix(self, seed=42):
        """生成低秩矩阵样本"""
        np.random.seed(seed)
        
        # 生成因子矩阵
        R1 = np.random.randn(self.n + self.n_prime, self.rank)
        R2 = np.random.randn(self.d + self.d_prime, self.rank)
        
        # 加权
        Sigma_diag = self.alpha ** np.arange(self.rank)
        X = R1 @ np.diag(Sigma_diag) @ R2.T / np.sqrt(self.rank)
        
        return X[:self.n, :], X[:self.n, self.d:],
               X[self.n:, :self.d], X[self.n:, self.d:]

不同规模下的表现

集中式设置：

方法	κ = 100	κ = 1000	κ = 10000	κ = 100000
梯度下降	$O(\kappa )$	$O(\kappa )$	$O(\kappa )$	$O(\kappa )$
共轭梯度	$O(\sqrt{\kappa })$	$O(\sqrt{\kappa })$	$O(\sqrt{\kappa })$	$O(\sqrt{\kappa })$
EAGLE	$O(\log \kappa )$	$O(\log \kappa )$	$O(\log \kappa )$	$O(\log \kappa )$

关键结果：在 $\kappa =10^4$ 时，EAGLE比共轭梯度法少约100倍的迭代次数。

分布式设置：

• 迭代复杂度与机器数量 $M$ 无关
• 收敛速度线性依赖于 ${\alpha }^{-1}$
• 在 $\alpha \approx 0.1$ 的低多样性场景中，加速EAGLE仍优于标准梯度下降约10倍

计算受限设置：

• 迭代次数随 $n/r$ 线性增长
• 每迭代时间随 $r$ 降低（最多7倍加速）
• 在 $r\ge n/4$ 时，EAGLE比随机梯度下降快约2.5倍

理论启示

上下文学习的本质是优化

这些发现揭示了上下文学习的深层机制：Transformer通过训练学会了执行数值优化算法。具体来说：

1. 数据空间 vs 参数空间：传统优化在参数空间迭代，而Transformer执行的是纯数据到数据的变换，没有显式的参数更新²。

2. 二阶动量的隐式出现：GD++和EAGLE中的预处理步骤实际上实现了二阶信息的使用，这与传统理解中Transformer仅实现一阶梯度下降的观点相悖⁷。

3. 算法发现的自动化：Transformer不仅学会了执行一个算法，还能在不同计算约束下自适应地调整该算法——这表明预训练可能”印刻”了一个灵活的基础算法框架。

隐式归纳偏置

为什么线性Transformer能够发现这些算法？

1. 架构约束：线性注意力的参数化自然地将更新分解为预处理和梯度步骤
2. 损失驱动：通过最小化MSE，Transformer被激励去学习能准确预测缺失数据的迭代规则
3. 对称性：权重的块对角结构反映了底层问题的数学结构

归纳偏置的证据：

• 每层仅需一个注意力头即可达到竞争性能
• 权重量化和稀疏化后算法性能保持稳定
• ${\alpha }_1^l{\alpha }_2^l\Vert A^l{\Vert }_2$ 在各层保持常数，表明Transformer学会了归一化批次最大谱范数²

参考资料

Footnotes

1. Brown, T., et al. (2020). Language models are few-shot learners. NeurIPS, 33, 1877-1901. ↩

2. Lutz, P., Gangrade, A., Daneshmand, H., & Saligrama, V. (2025). Linear Transformers Implicitly Discover Unified Numerical Algorithms. NeurIPS 2025. ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷ ↩⁸ ↩⁹ ↩¹⁰ ↩¹¹

3. Garg, S., et al. (2022). What can transformers learn in-context? A case study of simple function classes. NeurIPS, 35, 30583-30598. ↩

4. Akyürek, E., et al. (2023). What learning algorithm is in-context learning? Investigations with linear models. ICLR. ↩

5. von Oswald, J., et al. (2023). Transformers learn in-context by gradient descent. ICML, 35151-35174. ↩ ↩²

6. Hestenes, M. R., & Stiefel, E. (1952). Methods of conjugate gradients for solving linear systems. J. Res. Natl. Bur. Stand., 49(6), 409-436. ↩

7. Vladymyrov, M., et al. (2024). Linear transformers are versatile in-context learners. NeurIPS, 37, 48784-48809. ↩ ↩² ↩³