Rethinking Neural Nonlinearity as Gating¶

会议: NeurIPS 2026
arXiv: 2607.03148
OpenReview: https://openreview.net/forum?id=xc7nzeKG4m
领域: 优化/理论
关键词: 神经网络非线性、门控机制、阈值门控、极小分支定理、存算一体硬件

一句话总结¶

将传统线性变换后的激活函数重新表述为 MAC 运算前的输入条件阈值门控（Threshold Gating），在理论上证明主流激活函数及 Softmax 均为门控特例，并在无需重训的前提下实现跨 CNN、ViT、RNN 及大语言模型的无损转换与硬件统一。

研究背景与动机¶

现代深度神经网络的标准范式始终将层运算解构为「线性仿射变换 + 显式非线性激活函数」两步走机制，缺乏非线性激活的多层网络将直接退化为单层线性映射。数十年来，激活函数经历了从早期的 Sigmoid、Tanh 到分段线性的 ReLU 及其衍生变体（LeakyReLU、PReLU、ELU），再到平滑非单调函数（SiLU、GELU）的演化路线；然而在硬件底层视角下，这类离散架构选择为数字加速器与新兴模拟存算一体硬件（Analog In-Memory Computing, AIMC）带来了沉重的异构计算负担。在基于忆阻器（ReRAM/FeCAP）交叉阵列的模拟内存计算中，基尔霍夫电流定律与欧姆定律能以极低能耗和近乎零时延瞬时完成大规模矩阵乘累加（MAC），但芯片绝大部分面积和功耗却被模数/数模转换器（ADC/DAC）以及实现各类复杂非线性函数（如指数、误差函数、高斯积分等）的专用电路所吞噬。

传统激活函数的设计视角一直将非线性视为后 MAC 阶段的孤立非线性节点，这导致硬件必须为每一种激活函数维护专用算子或查表电路，且无法将非线性内化吸收至乘加运算中。更深层次的理论瓶颈在于：学术界尚未建立起一个能够统一定量刻画分段线性、平滑饱和、乃至跨元素 Softmax 注意力归一化的极小计算原语，模型设计者只能在预设的离散激活集合中做工程试错。

面对存算一体时代下从「优化 MAC 算力」向「消除转换与非线性电路瓶颈」的范式转变，本文重新审视了非线性产生的本质来源，提出激活函数的非线性完全可以等价转化为乘法前输入所驱动的分支阈值路由。核心 idea：将神经非线性解构为前置条件分支的阈值门控原语（Threshold Gating, TG），通过统一的可微分软门控插值与仿射分支，将所有主流激活函数及注意力 Softmax 统一为该原语的特例，建立极小分支定理并打通模拟交叉阵列的单一前置电路实现。

方法详解¶

整体框架¶

阈值门控（TG）彻底打破了「先求和后激活」的传统计算图边界。在数学表达上，常规全连接层 $a_j = \sigma(\sum_i W_{ji} x_i + b_j)$ 与将激活分配至输入端的表示 $a_j = \sum_i W_{ji} \sigma(x_i) + b_j$ 在多层级联管线中完全恒等。TG 进一步将单变量标量变换 $\sigma(x_i)$ 重构为基于 $K$ 个分支区间的门控聚合：每个输入分量 $x_i$ 首先驱动一组由 Sigmoid 构成的软阈值门控集合 $\{g_k(x_i)\}_{k=1}^K$，门控权重满足单位分解（$\sum_k g_k = 1$），并将分配权重作用于各分支的仿射映射 $s_k x_i + c_k$。

对于预训练网络的无损转换，TG 通过权重绑定机制（TGA 梯队）将分支参数直接关联到原始权重矩阵 $W_{ji}$，保持原有网络权值完全不变；而在端到端从头训练场景中，TG 则可释放为突触级或通道级分支参数，使网络自适应学习最契合自身任务表征的连续非线性曲率。该原语从前向数据流到硬件阵列的映射拓扑如下：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入向量 x_i"] --> B["阈值门控原语 (TG)<br/>计算软门控权重 g_k(x_i)"]
    B --> C["分支仿射计算与统一归纳<br/>加权聚合 y_i = Σ g_k (s_k x_i + c_k)"]
    C --> D["线性乘加变换 (MAC)<br/>阻抗交叉阵列或张量矩阵乘 a_j"]
    D --> E["跨元素归约扩展 (TG-Softmax)<br/>基于 LSE 动态阈值的二分支归一化"]

关键设计¶

1. 阈值门控通用原语与主流激活函数的参数化归一 传统激活函数在形态上看似千差万别（如硬折线的 ReLU 与平滑连续的 SiLU、GELU），但 TG 原语揭示了它们在代数结构上的内在统一。对于单输入标量 $x_i$，基础门控函数定义为含锐度参数 $\tau>0$ 与阈值 $\theta$ 的广义 Sigmoid 形式 $g(x_i; \tau, \theta) = \sigma(\tau(x_i - \theta))$。当扩展至具有 $K-1$ 个阈值 $\theta_1 < \dots < \theta_{K-1}$ 的 $K$ 个区域时，第 $k$ 个区间的门控 $g_k(x_i) \in [0, 1]$ 组合成单位分解。整个 TG 原语的形式化定义为： $$y_i = \sum_{k=1}^K g_k(x_i) (s_k x_i + c_k)$$ 当带入全连接层权重后，层输出可统一写为 $a_j = \sum_i \sum_{k=1}^K g_k(x_i) (w_{ji}^{(k)} x_i + c_{ji}^{(k)}) + b_j$。通过控制参数元组 $(K, \tau, \boldsymbol{\theta}, \mathbf{s}, \mathbf{c})$，现有激活函数均呈现为其精确或极高精度的退化解： - ReLU：属于乘法族（Multiplicative Family），取 $K=2, \tau \to \infty, \theta=0, \mathbf{s}=(1, 0), \mathbf{c}=(0, 0)$，由几乎处处收敛的性质 $\lim_{\tau \to \infty} \sigma(\tau x) = \mathbf{1}[x > 0]$ 实现严格无损等价； - SiLU：同样为乘法族，$K=2, \tau=1.0, \theta=0, \mathbf{s}=(1, 0), \mathbf{c}=(0, 0)$，软门控在代数上完全恒等于 $x \cdot \sigma(x)$； - Sigmoid / Tanh：属于常数项族（Constant Family，$\mathbf{s}=(0, 0)$），利用输入仅作为门控切换信号、分支输出固定偏置常数的机制，如 Tanh 严格对应 $K=2, \tau=2.0, \theta=0, \mathbf{c}=(+1, -1)$； - GELU：基于高斯误差积分近似，在标准 $K=2, \tau=1.702$ 下全局点对点误差 $\le 0.3\%$，若引入 $K=3$ 校准则误差收敛至 $0.04\%$。

2. 极小分支定理与激活复杂度驱动的可训练性 面对「多大规模的 $K$ 才能充分表达神经网络所需的非线性」这一本质问题，论文提出了极小分支定理（Minimal Branch Theorem）。定理指出，对于任意生产级网络中所采用的良态激活函数（导数有界且具有单一主拐点/极值点），仅需 $K=2$ 个分支即可在 $\pm \infty$ 渐近线上严格匹配原函数的渐近斜率 $s_{\pm}$ 并在中心阈值点保持连续过渡，而 $K=1$ 蜕变为纯仿射函数、无法提供有效非线性跃迁。为了度量拟合残差对优化动态的影响，定理引入了激活复杂度指标（Activation Complexity） $\kappa(f) := |f'''(\theta_f)|$（即中心拐点处的三阶导数绝对值）。局部 Taylor 展开表明，二阶逼近残差直接受控于 $\frac{1}{6}\kappa(f)|x - \theta_f|^3$。更为关键的是，该复杂度内在地统领了深层网络训练的稳定性与梯度退化边界：在二阶优化中，稳定步长界为 $\eta = \mathcal{O}(1/\sqrt{\kappa(f)})$；在混沌边缘（Edge of Chaos）均值场理论下，高斯预激活分布的二阶曲率集中于 $\mathbb{E}[\phi''(z)^2] \approx \kappa(f)^2(q + \theta_f^2)$，从而导出最大可训练网络深度与激活复杂度平方成反比： $$L_{\max} \propto \frac{1}{\kappa(f)^2}$$ 这在理论层面解释了为何 $K=2$ 足以支撑现代超深网络的稳健训练——因为过大的曲率 $\kappa$ 虽能增加局部拟合自由度，却会导致网络深度可训练性急剧恶化。

3. TG-Softmax 跨元素注意力等价转换 注意力机制中的 Softmax 涉及多维向量在分母上的指数求和，传统上被视为与逐点激活完全不同的运算类别。本文利用对数和指数（Log-Sum-Exp, LSE）代数恒等式，成功将自注意力归一化纳入 TG 原语范畴。对于输入 logit 向量 $z$，位置 $i$ 的预激活 Softmax 输出可写作关于其余元素排除集 $z_{-i}$ 的函数： $$\text{softmax}(z)_i = \sigma(z_i - \text{LSE}(z_{-i}))$$ 该结构天然对应一个常数族 $K=2$ 的 TG 实例：其分支斜率 $\mathbf{s}=(0, 0)$，分支常数 $\mathbf{c}=(1, 0)$，门控锐度 $\tau=1$，而阈值则由全局归约动态生成为 $\theta_i = \text{LSE}(z_{-i})$。进而，由于双变量 $\text{LSE}(a, b) = \max(a, b) + \text{softplus}(-|a - b|)$，其中 $\max$ 为严格 $K=2$ 的硬门控 TG，$\text{softplus}$ 则是软门控函数的直接积分，整个前向自注意力机制被完整解构为统一的 TG 门控操作配合模拟前端的一阶积分电路。

损失函数 / 训练策略¶

TG 原语既支持即插即用的推理期免微调转换，也支持端到端从头训练： - 推理无损转换：对于 ReLU、SiLU、Sigmoid、Tanh 及 Softmax，直接通过代数恒等封闭解替换 PyTorch 模块，浮点精度下模型输出完全无损（$\Delta = 0.00\%$）；针对 ViT 中使用的 GELU，采用 640 张未标注校准图像在预激活层进行逐层 RMS 最小化的 Nelder-Mead 优化，确定层级 $(\tau, \boldsymbol{\theta})$； - 梯度反向传播与可微分学习：在可训练软门控模式下，门控函数天然光滑可导，$\tau, \theta, s_k, c_k$ 均可接收原生梯度；即使在 $\tau \to \infty$ 的硬门控模式下，亦可借助次梯度逼近维持正常训练； - 分层参数控制梯队：根据参数开销定义了四级共享模式：General 原语（突触级独立分支，$3K\times$ 参数）、$\text{TG}_S$（吸收偏置，$2K\times$ 参数）、$\text{TG}_W$（通道/层级共享门控但保留 $K$ 个分支权重矩阵，$K\times$ 参数）、$\text{TG}_A$（绑定分支权重至原主权重，仅增加微量标量门控，额外参数量 $\le 1\%$）。

实验关键数据¶

主实验¶

为了验证 TG 替代传统激活函数的有效性，论文在两类场景下展开核心评测：第一是跨主流经典视觉模型与循环时序网络的免微调转换（ImageNet-1K / MNIST），第二是大语言模型端到端（含 FFN 激活与注意力 Softmax 全量替换）在 WikiText-2 和 OpenWebText 上的困惑度验证。

下表总结了典型架构的无损转换精度对比（ImageNet Top-1 准确率及语言模型 PPL）：

模型架构	原始激活配置	TG 分支数 $K$	原始 Baseline	转换后 TG 性能	性能变化 $\Delta$
ResNet-50	ReLU	2	80.35%	80.35%	0.00%
EfficientNet-B0	SiLU + Sigmoid	2	77.67%	77.67%	0.00%
ViT-B/16 (通用 $\tau$)	GELU	2	81.10%	80.01%	-1.09%
ViT-B/16 (逐层校准)	GELU	2	80.86%	80.46%	-0.40%
ViT-B/16 (校准 $K=3$)	GELU	3	85.14%	85.14%	0.00%
2-Layer LSTM	Sigmoid + Tanh	2	98.09%	98.09%	0.00%
2-Layer GRU	Sigmoid + Tanh	2	98.89%	98.89%	0.00%
SmolLM-135M	SiLU + Softmax	2	20.295 PPL	20.295 PPL	0.00%
TinyLlama-1.1B	SiLU + Softmax	2	10.160 PPL	10.160 PPL	0.00%
GPT-2 124M	GELU + Softmax	2	29.249 PPL	29.314 PPL	+0.22% (轻微漂移)

消融实验¶

在语音识别基准 AN4 上，针对 DeepSpeech2 网络的 5 层双向 LSTM 门控机制进行消融分析。该实验验证了激活函数「结构家族约束」的重要性：LSTM 的原生状态插值机制属于常数族（$s_k=0$），当违反常数族先验强行引入可学习的乘法斜率 $s_k$ 时，优化将破坏门控代数结构，导致词错误率（WER，越低越好）显著回退。

模型隐层维度 $H$	方案配置 / 门控结构	词错误率 (WER %)	相对 Baseline 表现
$H=128$	标准 LSTM 基线 (固定常数门控)	19.53	基线对照
$H=128$	$\text{TG}_A$ $K=2$ 乘法族配置 (违背家族约束)	16.82	降 2.71%
$H=128$	$\text{TG}_A$ $K=2$ 常数族保留配置 (本文提出)	15.01	大幅优化 4.52%
$H=256$	标准 LSTM 基线 (固定常数门控)	13.45	基线对照
$H=256$	$\text{TG}_A$ $K=2$ 乘法族配置 (违背家族约束)	13.58	劣于标准基线
$H=256$	$\text{TG}_A$ $K=2$ 常数族保留配置 (本文提出)	12.29	优化 1.16%
$H=512$	标准 LSTM 基线 (固定常数门控)	11.77	基线对照
$H=512$	$\text{TG}_A$ $K=2$ 乘法族配置 (违背家族约束)	12.42	明显掉点
$H=512$	$\text{TG}_A$ $K=2$ 常数族保留配置 (本文提出)	10.35	优化 1.42%

关键发现¶

训练收敛步数显著缩短：在采用通道级门控参数学习的 $\text{TG}_W$ 模式下，网络达到 Baseline 最终指标所需的训练 Step 大幅降低：在 CIFAR-10 上对标 SiLU 实现 $2.04\times$ 步数加速，CIFAR-100 上对标 GELU 达到 $1.98\times$，时序 DeepSpeech2 达到 $1.52\times$，而在 Mamba 混合状态空间模型（MambaMIM）上亦取得 $1.47\times$ 的步数加速（符号检验 $p \approx 4\times 10^{-9}$）。
网络宽度与深度压缩增益：借助 $\text{TG}_W$ 增强的非线性拟合表征，ResMLP 在通道膨胀比 $h=1$ 时即可达到标准 ReLU 模型在 $h=2\sim 3$ 时的分类精度，使中间激活张量大幅压缩，直接降低边缘计算中的访存搬运开销；在 ResNet 深度扫描中，$K=2$ 贡献了最陡峭的性价比提升曲线，而扩展至 $K=3$ 仅带来边际收益。
注意力与前馈解耦验证：在 LLM 的全量转换实验中，消融分项证明通过 LSE 动态阈值转换的 Softmax 在 SmolLM、TinyLlama 及 GPT-2 上均达到了 $0.00\%$ 的完全无损转换；GPT-2 出现的 $+0.22\%$ 困惑度微弱漂移完全来自于 GELU 的双分支截断近似，在改用 $K=3$ 标定后漂移完全消除。

亮点与洞察¶

将激活函数从孤立算子升华为通用分支门控原语：揭示了各类常用非线性在底层均是输入驱动阈值分支的特殊表现形式，以统一的代数视角打破了激活函数、多专家门控以及注意力的边界。
理论推导与硬件实现的严谨闭环：通过极小分支定理与激活复杂度严格限定了 $K=2$ 的充分性，并直接对齐模拟交叉阵列（RRAM/FeCAP）的前端差分对晶体管电路，从根源上缓解了模拟计算的 ADC/DAC 与非线性计算面积瓶颈。
全架构免重训的无损即插即用能力：不仅提供了面向新模型的自适应连续训练空间，还实现了主流大语言模型、视觉模型及 RNN 的零成本转换，展现出极高的通用工程价值。

局限与展望¶

实际墙上时间受限于缺乏专用算子优化：虽然在模型训练优化步数（Step）上取得了 $1.47\times \sim 2.04\times$ 的显著加速，但由于目前仅依靠 Triton 自定义算子实现，尚未针对现代 GPU 硬件底层深度流水定制，实际墙上时间与高度优化的传统 GEMM 相比尚未拉开量级差距。
模拟芯片级集成仍在拓扑验证阶段：文章提出了完整的 MOSFET 差分对与阻抗交叉阵列联合参考拓扑，但物理器件级的工艺波动、噪声干扰以及全片上系统的实际流片与硅验证仍有待后续工作推进。
仅适用于主拐点约束的良态激活：针对多波峰、多高频振荡的特殊周期性非线性函数（如某些物理信息神经网络 PINN 所需的周期激活），$K=2$ 甚至 $K=3$ 的门控分支在理论上不再适用。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 颠覆性地将神经非线性归纳为输入阈值门控原语，打通了传统激活、Softmax 与模拟存算硬件的鸿沟。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 CNN、ViT、RNN、Mamba 及主流开源 LLM，兼备零样本转换与从头训练，理论证明极为严密。
写作质量: ⭐⭐⭐⭐⭐ 概念铺垫自然，数学推演清晰优雅，理论与工程实现紧密呼应。
价值: ⭐⭐⭐⭐⭐ 对深度学习基础理论、低功耗边缘端模型压缩及新兴神经拟态硬件架构设计均具深远指导意义。

模型隐层维度 \(H\)	方案配置 / 门控结构	词错误率 (WER %)	相对 Baseline 表现
\(H=128\)	标准 LSTM 基线 (固定常数门控)	19.53	基线对照
\(H=128\)	\(\text{TG}_A\) \(K=2\) 乘法族配置 (违背家族约束)	16.82	降 2.71%
\(H=128\)	\(\text{TG}_A\) \(K=2\) 常数族保留配置 (本文提出)	15.01	大幅优化 4.52%
\(H=256\)	标准 LSTM 基线 (固定常数门控)	13.45	基线对照
\(H=256\)	\(\text{TG}_A\) \(K=2\) 乘法族配置 (违背家族约束)	13.58	劣于标准基线
\(H=256\)	\(\text{TG}_A\) \(K=2\) 常数族保留配置 (本文提出)	12.29	优化 1.16%
\(H=512\)	标准 LSTM 基线 (固定常数门控)	11.77	基线对照
\(H=512\)	\(\text{TG}_A\) \(K=2\) 乘法族配置 (违背家族约束)	12.42	明显掉点
\(H=512\)	\(\text{TG}_A\) \(K=2\) 常数族保留配置 (本文提出)	10.35	优化 1.42%