检索中心深度学习:增长型非参数神经网络 Retrieval-Centric Deep Learning in Growing Nonparametric Neural Networks
用核化与加权 softmax 注意力替换线性注意力,导出按需增长的键值记忆网络及其学习规则
前置知识
线性层的对偶形式(Dual Form)
梯度下降训练的线性层 $y=Wx$,其权重可展开为 $W_T = W_0 + \sum_{t=1}^T v_t x_t^\top$,于是任意输入的输出 $y = W_0 q + \sum_t v_t\, x_t^\top q$,等价于以训练输入 $x_t$ 为键、反向传播误差 $v_t = -\eta_t \nabla_{y_t}\mathcal{L}$ 为值、$q$ 为查询的未归一化注意力检索。参数化权重与增长式键值记忆在这种表示下逐点相等。
全文的出发点:RCDL 的全部目标就是把这个只对线性层精确成立的恒等式推广到更强的注意力核,并用它重新设计学习算法。
线性注意力(Linear Attention)
不经 softmax 归一化、直接以 $\phi(x_t)^\top q$(最简单即内积 $x_t^\top q$)加权值向量的注意力形式,因可写成 $\sum_t v_t\,\phi(x_t)^\top q$ 而能增量计算、状态恒定。序列建模研究普遍表明其表达能力劣于 softmax 注意力。
对偶形式中的检索恰好就是一种线性注意力,论文的核心问题正是:能否以及如何用更强的注意力函数替换它来改进深度学习。
正定核与再生核希尔伯特空间(RKHS)
满足对称性与半正定性的核函数,如 RBF 核 $\phi(x,x')=\exp(-\|x-x'\|^2)$。Moore-Aronszajn 定理保证任何正定核都可写成特征映射的内积 $\phi(x,x')=\langle\psi_\phi(x),\psi_\phi(x')\rangle_{\mathcal{H}_\phi}$,对应一个隐式的无限维特征空间,即经典核机器的数学基础。
论文 3.1 节据此把非线性核注意力层解释为在无限维权重矩阵 $W\in\mathbb{R}^{d_{out} imes\infty}$ 上做梯度下降,从而严格导出 key/value 写入规则。
softmax 注意力与温度
$s_D(q)=\frac{\sum_{(k,v)\in D}e^{\langle k,q\rangle}v}{\sum_{(k,v)\in D}e^{\langle k,q\rangle}}$,加温度时对内积缩放(本文用 $e^{\frac{T}{\sqrt{d_{in}}}\langle k,q\rangle}$)。分母的归一化使每条注意力权重依赖于数据库中全部键。
正是这个分母使新键的加入造成输出不连续跳变,是全文最核心的数学困难;论文为此设计加权 softmax,温度行为也是实验分析(Fig. 3)的重点。
泛函梯度下降(Functional Gradient Descent)
把可优化对象从参数向量推广为函数 $w(k,v)$,在函数空间(以高斯测度 $\mu$ 定义 $L^2(\mu)$ 内积)中用 Fréchet 导数与 Riesz 表示定理求出泛函梯度并做下降。对 $w$ 的函数加法更新恰好对应往数据库追加一条加权 key/value 记录。
加权 softmax 的学习规则(论文 Eq. 6-7)完全在这套泛函分析语言下推导,不熟悉它就读不懂 3.2 节的核心论证。
二阶优化与 KFAC/FOOF
利用损失曲率信息(如 Fisher 信息矩阵或输入协方差之逆 $C^{-1}$)对梯度做预条件的优化方法。KFAC 用块对角结构近似 Fisher 矩阵,FOOF 是其高效简化版,本质是对输入施加白化变换,接近自然梯度/牛顿法。
论文第 4 节证明 mesa-layer 的闭式解精确实现 FOOF,这是把序列建模的线性注意力进展与二阶优化谱系正式连接起来的关键结论。
研究动机
传统深度学习用梯度下降把任意规模的训练数据压缩进固定尺寸的权重矩阵 $W$。一个经典但少有人注意的性质是:任何线性层都有“对偶形式”——$W_T = W_0 + \sum_{t=1}^T v_t x_t^\top$ 可精确改写为一个存储全部训练样本 (key, value) 对、用未归一化点积注意力(即线性注意力)做检索的键值记忆系统,其中键 $k_t=x_t$、值 $v_t=-\eta_t\nabla_{y_t}\mathcal{L}$。Irie 等人据此提出自然的问题:既然线性注意力在序列建模中已被证明明显劣于 softmax 注意力,能否用更强的注意力函数替换这一底层机制来获得更强的深度学习系统?这个问题一直悬而未决,原因是直接把 softmax “插入”线性注意力会破坏原有 key/value 的最优性:softmax 分母涉及对全部键的归一化,新键哪怕带极小的值向量也会让输出不连续地跳变,朴素套用线性情形的学习规则不再对应任何有原则的优化过程;同时非线性核下模型每个训练步都在增长,内存与计算如何控制也亟待回答。
本文的目标是本文要把“检索中心”的直觉变成一个理论基础扎实、可以真正训练的新范式,具体目标有三。其一,为非线性核化的注意力层建立严格的函数梯度学习规则:先处理正定核(如经典的 RBF 核),证明对隐式 Hilbert 空间中无限维权重矩阵做梯度下降在数学上恰好等价于每步往数据库追加一对 key/value;再处理更难的 softmax——作者专门设计“加权 softmax 核”并推导其学习算法,使系统的每一步更新都在最小化给定全局损失,从而正式确立“检索中心深度学习”(RCDL)范式。其二,把 MesaNet 与 DeltaNet 这类先进线性注意力替换进对偶形式,证明它们摇身一变成为传统固定尺寸网络的新型优化器,并与近期提出的 FOOF 优化器及二阶优化建立形式联系。其三,针对模型线性增长的核心痛点,提出混合学习算法:仅在输入与已有键足够不同时才增长数据库,把内存增长压到次线性。
与已有工作不同的是,论文的独特切入角度有三层。第一,粒度上的“逐层精确对偶”:与 Domingos 等人把整个训练好的深度网络近似看作一个核机器不同,RCDL 建立在“每个由梯度下降训练的线性层都精确等价于一个(线性)核机器”这一逐层恒等式上,且各层特征仍由端到端梯度下降训练——深度学习的核心原则得以保留,这为研究网络优化开辟新途径。第二,提问方向的翻转:不从“如何更新固定权重”出发,而是反过来问“该往数据库写入什么 key/value 才等价于对全局损失做优化”,把学习算法的设计问题转化为检索库的写入规则问题。第三,工具上的跨界:为驯服 softmax 归一化的不连续性,作者引入高斯混合密度、Radon-Nikodym 导数、Fréchet 导数与 Riesz 表示定理这套泛函分析工具,在加权 softmax 函数构成的空间里做泛函梯度下降——这种处理方式在深度学习文献中相当罕见,却意外给出了闭式解。
核心方法
直觉上,参数网络把记忆压缩进固定权重矩阵,RCDL 则把每个训练样本显式写成一条可检索记录,检索质量决定学习质量。技术路线分三步。第一步(正定核):由 Moore-Aronszajn 定理,正定核 $\phi(x,x')=\langle\psi_\phi(x),\psi_\phi(x')\rangle_{\mathcal{H}_\phi}$ 使非参数层可视为隐式 Hilbert 空间上的无限维权重矩阵 $W\in\mathbb{R}^{d_{out}\times\infty}$,其梯度更新 $\Delta W\propto(-\eta_t\nabla_{y_t}\mathcal{L})\,\psi_\phi(x_t)^\top$ 恰好等价于追加 $(k_t,v_t)=(x_t,-\eta_t\nabla_{y_t}\mathcal{L})$,RBF 核即属此类。第二步(加权 softmax):为避免标准 softmax 分母跳变,给每条记录配非负标量权重 $w$,检索函数 $s_D(q)=\frac{\sum_{(w,k,v)\in D}w\,e^{\langle k,q\rangle}v}{\sum_{(w,k,v)\in D}w\,e^{\langle k,q\rangle}}$;把数据库解释为 $(k,v)$ 空间上的高斯混合密度后,对其诱导的连续 softmax 函数做泛函梯度下降,得到闭式写入规则。第三步(高级线性注意力):把 mesa-layer(正则化最小二乘意义下的最优线性映射)与 DeltaNet(delta 规则的在线梯度步)插入对偶形式,mesa-layer 的闭式解精确实现 FOOF 二阶优化器,DeltaNet 给出与 Muon 竞争的实用优化器。
核心创新是“加权 softmax 核”及其泛函梯度学习规则。为什么直接把 softmax 塞进线性注意力行不通?因为 softmax 分母 $\sum_{(k,v)\in D}e^{\langle k,q\rangle}$ 依赖全部键,任何新键的加入——即使值向量为零——都会不连续地改变输出,因此不存在“无穷小地添加一条记录”意义上的梯度。作者的办法是把每条记录升级为带权重 $w$ 的三元组 $(w,k,v)$。妙处在于:把有限加权数据库看作 $(k,v)$ 联合空间上的高斯混合密度 $w_D(k,v)=\sum_i w_i\,\mathcal{N}\big((k,v);\,(k_i,v_i),\,\Sigma\big)$,其中 $\Sigma=\mathrm{diag}(\sigma_k^2 I_{d_{in}},\sigma_v^2 I_{d_{out}})$,代入连续 softmax 后方差项 $e^{\frac{1}{2}\sigma_k^2\|q\|^2}$ 在分子分母中正好抵消,离散加权 softmax 与连续函数严格相等。于是“往数据库加一条记录”等价于“在函数空间加一个指数基函数”,泛函梯度 $\nabla_w J(k,v)=\frac{e^{\langle k,q_t\rangle}}{D(w)}\langle\hat\nabla_y\mathcal{L},\,v-s_w(q_t)\rangle$ 就有闭式解,且新键 $k_{new}=\Sigma_k q_t=\sigma_k^2 q_t$ 揭示 $\sigma_k^2$ 扮演逆温度角色。与已有方法的本质区别:学习规则不再是对固定参数的更新,而是对不断生长的函数空间中测度密度的泛函梯度。
方法步骤详情
以残差 MLP 块 $e\leftarrow e+\sigma(W_2\sigma(W_1\mathrm{RMSNorm}(e)))$ 为骨干,把每个线性块换成增长层 $e\leftarrow e+g_D(\mathrm{RMSNorm}(e))$,$g_D$ 取加权 softmax 层或 RBF 层,输出层为分类用的 softmax 层或回归用的线性层。训练循环如下。前向:对查询 $q$ 用带温度的加权 softmax $s_D(q)=\frac{\sum e^{\frac{T}{\sqrt{d_{in}}}\langle k,q\rangle}\tilde v}{\sum we^{\frac{T}{\sqrt{d_{in}}}\langle k,q\rangle}}$ 计算输出,存储加权值 $\tilde v=wv$ 保证数值稳定。写入:对每个观测 $(q_t,y_t)$ 计算预测 $s_t=s_{w_t}(q_t)$ 与误差信号 $g_t=\hat\nabla_y\mathcal{L}$,追加三元组——权重 $w_{new}\propto\eta\langle s_t,g_t\rangle$、键 $k_{new}=q_t$(配合温度 $T$)、加权值 $\tilde v_{new}\propto-\eta\sigma_v^2 g_t$;mini-batch 时利用梯度的线性性,一步同时追加 $B$ 条记录。负权重处理:$w_{new}<0$ 会导致分母相消,可用截断 $w_{new}\leftarrow\max(0,w_{new})$,或利用规范不变性(值整体平移 $+b$ 则输出平移 $+b$)选虚拟规范 $b_t$ 保证投影 $\langle g_t,s_t+b_t\rangle>0$。混合模式:把已有 $(k,v)$ 当作可学习参数交给 Adam/AdamW 微调,仅当 $\min_j\|x-k_j\|_2>\tau$ 时才增长数据库,实现次线性内存。优化器方向则令 $v_t=W_{t-1}x_t-\eta_t\nabla_{y_t}\mathcal{L}$,mesa-layer 闭式解 $(V_{<T}X_{<T}^\top+V_TX_T^\top)(X_{<T}X_{<T}^\top+X_TX_T^\top+\lambda I)^{-1}$ 直接给出 FOOF 更新规则。
技术新颖性
新颖性有四点。第一,填补理论空白:Irie 等人提出“用更强注意力替换对偶形式中的线性注意力”后,本文首次严格证明对正定核增长网络,无限维权重矩阵上的梯度下降与追加 key-value 严格等价,把启发式想法变成有原则的学习算法。第二,加权 softmax 核是为解决归一化不连续性量身设计的新对象,配套推导(高斯混合诱导的连续 softmax、经 Radon-Nikodym 导数换测度、Fréchet 导数加 Riesz 表示求泛函梯度)在深度学习中独树一帜,作者也指出其精神与泛函梯度提升一脉相承,并与早期用函数梯度逐模块构建网络的工作相衔接。第三,规范不变性是处理负权重的优雅机制:值向量整体平移 $+b$ 输出随之平移 $+b$,同一函数在不同“规范”下有不同表示,动态选规范即可保证 $w_{new}>0$,比简单截断更有原则。第四,跨领域的形式联系:证明 mesa-layer 在对偶形式下的闭式解精确实现 FOOF——一个由 KFAC 简化而来的二阶优化器——其加速来自额外的逆矩阵对输入做白化,从而把序列建模社区的 MesaNet/DeltaNet 进展与经典二阶优化谱系形式地连接起来。
实验结果
合成任务:教师是 100 维输入、单隐层 ReLU、残差流宽 100、LeCun 正态初始化且权重固定的随机 MLP,学生用在线 SGD 训练;过参数化的参考 MLP 最强作为下界,但仅有一个可训练层的 GWS 显著优于宽度与最终数据库规模匹配、用 Adam 训练的参数化 softmax 基线。CIFAR-10/100:GWS 与 RBF 增长模型仅用普通 SGD 就显著超过 SGD/Adam/AdamW/RMSProp 训练的残差 MLP,与 Muon 相当,GWS 略胜 RBF;Mesa 衍生的 FOOF 在 CIFAR-100 上全场最佳,Muon 在 CIFAR-10 上最佳;增长模型训练损失下降远快于所有基线且保持良好泛化。单遍在线学习(1 epoch、无数据增强):增长方法超过包括 Muon 和 Mesa 在内的全部参数模型。分析实验:深度增加性能提升但约 2 个隐层饱和;温度呈非单调依赖,最优 $T=10$,在 $d_{in}=3072$ 下 $T/\sqrt{d_{in}}\langle k,q\rangle\in[-5.54,5.54]$,说明层工作在非线性区——训练后把 exp 线性化导致性能骤降,证实模型真的在用非线性核。对比实验:参数化 softmax 基线训 2000 epochs 出现持续扩大的训练-测试差距且调参无法消除,GWS 训 50 epochs 无此问题。ViT 试验:把每层 FFN 替换为 RCDL 层后,CIFAR-10 单 epoch 从纯参数 ViT 的 48.4% 提高到 54.4%。混合层:以阈值 $\tau$ 控制增长,在保留高精度的同时实现次线性内存,进一步压缩时性能平缓退化。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Teacher-student 合成回归(100 维输入) | 回归损失(mean ± s.e.m.) | GWS 单层增长模型,损失显著低于等容量参数化 softmax 基线(Fig. 1) | 宽度匹配最终样本数、Adam 训练的参数化 softmax MLP | 在仅一个可训练层的架构劣势下仍显著更低(论文未给出绝对数值) |
| CIFAR-10 图像分类 | 测试精度 | GWS/RBF(plain SGD)显著超过 SGD/Adam/AdamW/RMSProp MLP,与 Muon 相当;单遍学习中最优 | Muon 训练的残差 MLP(CIFAR-10 最佳参数基线)、Mesa/FOOF | 相对同 SGD 基线大幅领先;单遍设置下超过全部参数模型 |
| CIFAR-100 图像分类 | 测试精度 | GWS/RBF 超过 SGD/Adam 系基线 | Mesa/FOOF 优化器(该数据集全场最佳) | 增长模型未胜过 FOOF,凸显为增长模型开发高级学习规则的必要性 |
| CIFAR-10 单 epoch(ViT 中替换 FFN) | 测试精度 | RCDL 混合 ViT 54.4% | 纯参数 ViT 48.4% | +6.0 个百分点 |
| CIFAR-10 混合增长(1 epoch) | 测试精度 vs 最终容量 | 高阈值 τ 下保持高精度、次线性内存;压缩时平缓退化(Fig. 5) | 全量非参数增长(线性内存) | 以远小于全量的容量逼近全量性能 |
局限与改进
作者承认的首要限制是内存与计算随训练步数线性增长 $O(T)$:选非线性核后每个训练步都要追加记录,大数据集上不可行;混合层只是初步缓解,有原则的压缩/剪枝方案仍是开放问题(transformer 文献的滑动窗口、token merging 或可借鉴)。其次,理论目前只支持基于普通 SGD 的学习规则(附录 D 扩展了权重衰减与动量),尚不知道如何为 RCDL 推导 Adam、Muon 这类自适应优化器。我的补充观察:实验规模局限于 CIFAR 级别的 MLP 与小型 ViT,没有语言模型等大规模验证;推理时检索成本随数据库增长,服务延迟与索引问题未讨论;负权重的截断与规范技巧是启发式处理,缺少收敛性分析;混合层把旧 $(k,v)$ 当参数微调可能覆盖早期记忆,作者未研究其灾难遗忘行为;此外增长模型用 SGD 而基线扫遍多种优化器,虽然作者做了密集超参搜索,但两类方法的搜索空间不同,公平性仍有讨论余地。
独立分析的弱点
第一,可扩展性是硬伤:CIFAR-10 五万样本乘以多个 epoch 意味着百万级记录,即便值向量仅几十维也达 GB 级内存,ImageNet 或语言建模完全不可行;改进方向是乘积量化、近似最近邻检索与有原则的键值合并。第二,增长判据 $\min_j\|x-k_j\|_2>\tau$ 需要扫描全部已有键,每次 $O(T)$,与次线性增长的目标自相矛盾;可用局部敏感哈希或层次化索引加速。第三,超参敏感:温度 $T$ 与测试精度呈非单调关系(Fig. 3 左),$\sigma_k^2$、$\tau$ 等也需按任务调节,实际部署成本高;可研究自适应温度或基于输入分布的自动调节。第四,学习规则受限于 plain SGD,无法享受自适应学习率收益,向 Adam 类规则的理论扩展是明确缺口。第五,规范不变性方法只在附录 G.2 给出初步结果,主实验多用截断,而截断会系统性丢弃负梯度方向的信息,值得更深入研究。
未来方向
作者明确提出的方向:压缩增长的非参数层(有原则的键值剪枝与合并)、把 RCDL 层集成进语言模型、为增长模型导出 Adam/Muon 级别的学习规则、以及面向持续学习的应用——RCDL 保留了过去的激活,未来可加入在前向传播中“复活旧记忆”的算法,这与神经科学中遗忘源于检索失败而非记忆丢失的“沉默记忆”发现相呼应。基于本文成果可延伸的:一是用 RCDL 层做在线适应与上下文记忆,替代 LoRA 等参数高效微调,实现无需重写权重的持续学习;二是利用 Fig. 5 的容量-性能曲线做弹性推理,按内存预算动态决定数据库规模;三是把规范不变性技巧推广到其他归一化注意力与概率模型;四是深挖 mesa-layer 与 FOOF 的等价性,用曲率预条件思想改造 DeltaNet 系列优化器;五是研究混合层容量分配的理论化,把增长判据与数据分布理解结合。
复现评估
论文提供了公开代码:https://github.com/paradigms-of-intelligence/rcdl。实验全部基于公开数据集 CIFAR-10/100 与可合成的 teacher-student 任务,无需专有数据。算力方面:骨干是残差 MLP 与小型 ViT,结果为 5 个种子平均,主实验单卡即可复现;需要注意的是增长模型的内存随训练样本线性增长,CIFAR 规模尚可(GWS 训 50 epochs),更大数据集需配合混合层。关键超参在文中有明确数值或范围:最优温度 $T=10$、输入维度 $d_{in}=3072$、增长阈值 $\tau$(Fig. 5 扫描多档)、正则化 $\lambda$;超参搜索协议在附录 F.1,混合层细节在 F.4,规范不变性版本在 G.2,高斯混合推导在附录 B/C。总体复现难度中等偏上:代码与数据可得、公式推导完整,但泛函梯度部分数学较深,密集超参扫描也需要一定计算预算。
论文图表