可逆神经网络(Reversible Neural Network,常写作 RevNet 或「可逆残差」)是一类把网络层设计成可逆变换的模型:前向算完之后不保存中间激活(activation),直接由输出反推出输入,从而把训练显存压到几乎与网络深度无关。
它为什么能省显存
普通 ResNet 每一层的输出 y = x + F(x) 都要留在显存里,等反向传播时再取用。网络越深,存得越多,显存随深度线性增长。
可逆残差的做法是把输入通道劈成两半 x₁、x₂,改成交叉更新:
y₁ = x₁ + F(x₂)
y₂ = x₂ + G(y₁)
这个变换没有丢失信息——想知道原来的 x₁、x₂,倒着减一遍就行:
x₂ = y₂ − G(y₁)
x₁ = y₁ − F(x₂)
打个比方:两个人合写一本账,A 把自己那页原样保留,B 的新数字是「旧数字 + 根据 A 算出来的量」。事后按相反顺序把这一项减掉,旧账就能完整还原。所以训练时只要留下每一段最后的输出,反向传播走到哪一段,就现场反推出这一段里的所有中间值,不必提前囤着。
和梯度检查点的区别
梯度检查点(AI 词典:Gradient Checkpointing">Gradient Checkpointing)是另一条思路:前向只存少量「检查点」,反向时从最近的检查点重跑一遍把缺的激活算出来。两者都是拿计算换显存,但性格不同。
| 可逆网络 | 梯度检查点 | |
|---|---|---|
| 省显存靠 | 由输出反推输入 | 只存检查点,其余重算 |
| 对模型要求 | 结构必须可逆,层要重写 | 任意模型,外面包一层即可 |
| 权衡粒度 | 基本只有「全可逆」一档 | 存几个点,连续可调 |
| 数值风险 | 反推放大误差,低精度下易漂 | 重算走原路,数值一致 |
| 显存量级 | 理论上与深度无关 | 最优可达深度平方根级 |
为什么输给了梯度检查点
不是思路不好,而是工程上不划算。可逆网络要你重新设计每一层:拆通道、改残差、处理下采样和归一化,还要保证整条链路都能倒推;切到半精度训练时,反推的误差累积还可能造成数值漂移。梯度检查点只是训练时的一个开关,模型代码几乎不用动,什么网络都能套,代价是多跑一遍前向(训练计算量大概多三成)。而省显存的效果已经够用——平方根级别的显存就能训下绝大多数模型。于是社区选了更通用、更好接、更可控的那个。
对从业者的意义
显存不够时,先开梯度检查点,这是最省事的一步;还不行再考虑并行、卸载到 CPU、量化等手段。可逆结构真正的舞台在别处:归一化流(Normalizing Flow)必须靠可逆性计算概率密度,Reformer 这类长序列模型用它来压注意力层的显存。值得记住的核心是一句话——激活值不必都存着,能重算出来就够了。
各框架对梯度检查点的支持方式和参数名不尽相同,具体以官方文档为准。
