面向不变架构的精确位级可逆 Transformer

各种可逆深度神经网络（DNN）模型已提出以减少训练过程中的内存消耗。然而，几乎所有现有的可逆 DNN 要么需要特殊的非标准架构，要么是通过显著修改现有 DNN 架构来实现可逆性。本文提出了 BDIA-transformer，这是一个在推理过程中使用不变标准架构的确切位级可逆 transformer。其基本思想是将每个 transformer block 视为求解常微分方程（ODE）的 Euler 积分近似，然后将双向积分近似（BDIA）技术整合到神经网络架构中，并结合激活量化以实现确切的位级可逆性。在训练过程中，我们让 BDIA-transformer 中的一个超参数 $\gamma$ 在每个训练样本的每个 transformer block 中随机取 $\{0.5, -0.5\}$ 两个值中的一个，以对每两个连续积分近似进行平均。结果表明，BDIA-transformer 可被视为通过一组二进制随机变量参数化的 ODE 求解器的训练，从而正则化模型并提高验证准确率。在前向过程中，需要存储每个 transformer block 的轻型副信息，以 account for 二进制量化损失以实现确切的位级可逆性。在推理过程中，取期望 $\mathbb{E}(\gamma)=0$ ，以使 BDIA-transformer 的结果架构与 transformer 相同，仅在激活量化方面有所不同。我们在图像分类和语言翻译上的实验表明，BDIA-transformer 在验证性能方面显著优于其传统 counterparts，同时还大幅降低了训练内存需求。