【DL輪読会】Reversible Deep Equilibrium Models

DEEP LEARNING JP [DL Papers] Reversible Deep Equilibrium Models Presenter: Shinji Kotani http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 • 論文名 – Reversible Deep Equilibrium Model (2025 9/16) • 著者 – Sam McCallum, Kamran Arora, James Foster • 発表学会 – 情報なし プレプリントのみ • リンク – https://arxiv.org/abs/2509.12917 2

• https://arxiv.org/abs/2509.12917

概要 Deep Equilibrium Models (DEQ)は少ないメモリで学習出来て、GPTと 同等のPerplexity(PPL)を達成するが、推論、学習速度が遅い。 これは学習で近似的な勾配を計算するからである。 Reversible Deep Equilibrium Model (Rev DEQ)では、陰関数定理を用 いず、正確な勾配を計算する可逆ソルバーを提案する。 • 学習におけるメモリ効率の維持 • Transformer-XL、DEQと比較して低いPPL • 収束が早くなることによる推論、学習速度の上昇 3

先行研究 Transformer Deep Equilibrium Models (DEQ) (2019 10/28) RNNでは経験的に、繰り返しを続けると 状態が一定に収束する。 DEQでは、RNNと同様に一つの層を何度 も繰り返し、その状態が収束した時に推 論を終了とする。 This is a embedding DEQ-Transformer This is embedding Transformer-block 1 繰り返し Transformer-block 2 Transformer-block 3 Transformer-block (DEQ ver) もし、状態zが 収束していれば Transformer-block 12 収束させるため、繰り返している間は入 力は固定である必要がある。 DEQをTransformer-XLに対して適用したも のがDEQ-Transformerである。 a head head is is a pen a pen DEQ-Transformerは一つの層を 何度も繰り返す 5

DEQ-Transformerの再帰処理 • 入力注入 入力を追加し続けることで、入力文 章を忘れないようにする。 • KVキャッシュは使えない。 • Solverが 𝒛𝒕 , 𝒚 から次の状態を決める。 • Solverにはブロイデン法、アンダー ソン法(主流)などがある。 状態 zt 入力 x linear linear Q,K,V 残差 接続 + Scaled Dot-Product Attention Feed Forward Network 残差 接続 出力 y 𝑧𝑡+1 = Solver(𝑧𝑡 , 𝑦) 6

陰関数定理 • 収束したと仮定すると、陰関数定理で 損失関数に対する勾配を計算できる。 • ただし巨大なヤコビアンの逆行列を求 めることは難しいため、ブロイデン法 ∂L ソルバーなどでw = ∗ (Jgθ −1 |z ∗ ) を ∂z 近似的に求める。 • そのため、近似的な勾配となる。 • BPTTを使わないため学習における、 メモリ効率が高い RNN zt+1 = f (zt , x) 収束した (zt+1 = zt ) と仮定 DEQ z ∗ = f (z ∗ , x) 陰関数定理を用いた 勾配計算 ∗ ∂L ∂L −1 ∗ ∂fθ (z ; x) = − ∗ (Jgθ |z ) ∂θ ∂z ∂θ DEQ層 より前 の勾配 ヤコビアン の逆行列 w DEQ層 より後 の勾配 8

DEQの実験結果 DEQ-Transformerは、同じパラメータで Transformer, Transformer-XLと比較して、同等の Perplexityを達成する。 また、陰関数定理により、学習でメモリ消費を 大きく削減できる。 ただし、状態が収束するまで繰り返しを行うた め、推論、学習速度が遅い。 学習で省メモリだが 推論、学習速度が遅い Transformer 9

DEQの問題点 陰関数定理 BPTT 良い点 良い点 • 学習のメモリ効率がとてもいい • 正確な勾配 • BPTTにおける繰り返し数を無限にした と言える。 • 入力を固定する必要がない 悪い点 悪い点 • 陰関数定理による近似的な勾配 • 勾配計算で収束を仮定しているが、厳 密には収束しないため、さらに近似的 な勾配となる。 二重の近似がある • 繰り返し数を増やすとメモリが線形で 増加 • 勾配爆発、消失 10

Reversible Deep Equilibrium Model (Rev DEQ) DEQにおいての、近似的な勾配は陰関数定理が原因なので、 陰関数定理を使わない！ 代わりに可逆ソルバーを使うことでBPTTのメモリ問題と、 勾配爆発収束問題を解決する DEQをBPTTで学習させる。 さらにそのBPTTを改良 11

可逆ソルバー 順伝播 可逆ソルバーでは二つの状態zt , yt を持つ。 yt+1 = (1 − 𝛽)yt + 𝛽f𝜃 (zt , x) zt+1 = (1 − 𝛽)zt + 𝛽f𝜃 (yt+1 , x) BPTTでは途中の状態をメモリに保存するた めメモリ効率が悪い しかし、可逆ソルバーを用いると、現在の状 態zn+1 , yn+1 から、zt , yt を厳密に復元出来る。 これにより、メモリに保存する必要がなくな り、モデルの実質的な深さNに対して𝑶(𝟏)と なる。 状態の復元 zt+1 − 𝛽f𝜃 (yt+1, x) zt = 1−𝛽 yt = yt+1 − 𝛽f𝜃 (zt , x) 1−𝛽 12

実験結果 言語モデリングタスク • 同パラメータでTransformer-XL, DEQに 比べて低いPerplexity(PPL) • DEQに比べて、少ない関数評価 （推論、学習速度3倍） 画像認識タスク • 少ないパラメータ数で同等の精度 • DEQに比べて、少ない関数評価 （推論、学習速度2~3倍) 13

まとめ・感想 まとめ：DEQのメモリ効率を維持し、正確な勾配を求め精度、推論学習速度を上昇 良い点 • 精度、推論学習速度を上昇している。 • 理論、数式がシンプルで理解しやすい 疑問が残った点 • 169Mまでしか検証されていない。スケーリングはするのか？ • PPLしか比較していない。下流タスクの性能は？ • BPTTとの学習速度、メモリ効率の比較がない。 14

参考文献 1. Bai, S., Kolter, J. Z., & Koltun, V. (2019). Deep equilibrium models. Advances in neural information processing systems , 32. 2. McCallum, S., Arora, K., & Foster, J. (2025). Reversible Deep Equilibrium Models. arXiv preprint arXiv:2509.12917. 15