跳到主内容
快讯直播
AI智模界
AI 词典

可逆神经网络:把中间激活反推出来

可逆神经网络(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 这类长序列模型用它来压注意力层的显存。值得记住的核心是一句话——激活值不必都存着,能重算出来就够了。

各框架对梯度检查点的支持方式和参数名不尽相同,具体以官方文档为准。

AI 生成本文由 AI 基于公开信息自动生成,仅供参考。