[论文解读] Deep Equilibrium Models

Shaojie Bai, J. Zico Kolter+1 more|arXiv (Cornell University)|Sep 3, 2019

Topic Modeling被引用 245

一句话总结

引入 Deep Equilibrium Model (DEQ)，它通过寻求权重绑定的非线性变换的固定点来建模具有恒定内存的序列，并在语言建模上展示具有显著内存节省的竞争性性能。

ABSTRACT

We present a new approach to modeling sequential data: the deep equilibrium model (DEQ). Motivated by an observation that the hidden layers of many existing deep sequence models converge towards some fixed point, we propose the DEQ approach that directly finds these equilibrium points via root-finding. Such a method is equivalent to running an infinite depth (weight-tied) feedforward network, but has the notable advantage that we can analytically backpropagate through the equilibrium point using implicit differentiation. Using this approach, training and prediction in these networks require only constant memory, regardless of the effective "depth" of the network. We demonstrate how DEQs can be applied to two state-of-the-art deep sequence models: self-attention transformers and trellis networks. On large-scale language modeling tasks, such as the WikiText-103 benchmark, we show that DEQs 1) often improve performance over these state-of-the-art models (for similar parameter counts); 2) have similar computational requirements to existing models; and 3) vastly reduce memory consumption (often the bottleneck for training large sequence models), demonstrating an up-to 88% memory reduction in our experiments. The code is available at https://github.com/locuslab/deq .

研究动机与目标

• 通过寻求权重绑定变换的平衡点，提出一种内存高效的深序列建模方法。
• 提出一个通用的 DEQ 框架，直接求解序列级固定点并对其进行微分。
• 展示 DEQ 在 TrellisNet 和权重绑定的 Transformer 上的大规模语言任务中的实例化。

提出的方法

• 定义平衡条件 z* = fθ(z*; x) 并使用黑盒根查找方法求解 z*。
• 通过隐式微分（定理1）对平衡点进行反向传播，避免存储中间激活。
• 使用布罗伊登的拟牛顿更新在前向和后向传播中近似逆雅可比，加速训练。
• 展示普适性：堆叠多个 DEQ 不会超过单个 DEQ 的表示能力（定理2）。
• 用 TrellisNet（权重绑定的 TCN）和具记忆增益的自注意力 Transformer 来实例化 DEQ，以覆盖主流序列模型家族。

实验结果

研究问题

• RQ1 权重绑定的深度序列模型的固定点形式是否能在语言建模任务中达到或超过层级堆叠结构的性能？
• RQ2 通过平衡点的隐式微分，是否可以实现 DEQ 模型的常量内存反向传播？
• RQ3 DEQ 在 TrellisNet 和 Transformer 上的实际实例化是什么，在 WikiText-103 和 PTB 等基准上表现如何？
• RQ4 DEQ 的内存占用和训练效率与传统深度网络及梯度检查点比较如何？
• RQ5 是否需要或有益于多层 DEQ 的堆叠以获得额外的表示能力？

主要发现

• DEQ 在 PTB 和 WikiText-103 上的困惑度可达到与同等规模的最先进模型竞争，甚至更好。
• DEQ 在训练过程相对于基于层的等价物节省超过 80% 的内存（最高可达 88%），并且在内存使用上可能优于梯度检查点。
• 通过隐式固定点公式计算反向梯度，使反向传播可实现常量内存。
• 两种实例化，DEQ-TrellisNet 与 DEQ-Transformer，证明了与卷积网络、RNN、Transformer 架构的兼容性，同时使用的资源相似或更少。
• 前向和后向在拟牛顿方法下收敛到平衡点，运行时比固定深度模型略高，但内存效率显著提高。