深度均衡模型（Deep Equilibrium Models）

（2.0 版本现已发布！:grinning:）

新闻

:boom:2021/6：仓库已更新，包含多尺度 DEQ（MDEQ）代码、Jacobian（雅可比矩阵）相关分析与正则化支持，以及通过 PyTorch 的 backward hook 实现的全新、更快且更简洁的隐式微分（implicit differentiation）方法！（参见 此处。）

• 如果你想从一个简化版的 DEQ 入手，NeurIPS 2020 关于“深度隐式层（Deep Implicit Layers）”的教程提供了详细的逐步介绍：教程视频与 Colab 笔记本在此。

• 一个基于 JAX 的 DEQ 实现（包括 Broyden 方法等的 JAX 版本）可在 此处 获取。

---

本仓库包含了深度均衡（Deep Equilibrium, DEQ）模型的代码。DEQ 是一种隐式深度（implicit-depth）架构，它直接求解并反向传播通过（有效意义上）无限深网络的（不动点）均衡状态。重要的是，与先前的隐式深度方法（例如基于 ODE 的方法）相比，本工作还展示了这种隐式模型与现代结构化层（如 Transformer）的强大潜力和兼容性，使得 DEQ 网络在不使用“深度”堆叠的情况下（因此内存复杂度为 O(1)），即可在 NLP 和视觉任务上达到与当前最先进（SOTA）深度网络相当的结果。此外，我们还提供了用于正则化这些隐式模型稳定性的工具。

具体而言，本仓库包含以下论文的代码（参考本文末尾的 bibtex）：

• Deep Equilibrium Models
• Multiscale Deep Equilibrium Models
• Stabilizing Equilibrium Models by Jacobian Regularization.

前置要求

Python >= 3.6 且 PyTorch >= 1.10。强烈建议使用 4 块 GPU 以获得计算效率。

数据

我们在 DEQ-Sequence/ 和 MDEQ-Vision/ 子文件夹中提供了更详细的数据集（如 WikiText-103、ImageNet、Cityscapes 等）下载与处理说明。

如何构建/训练一个 DEQ 模型？

自 2021 年 6 月起，我们将仓库划分为两个部分，分别包含序列模型 DEQ（即 DEQ-Sequence/）和视觉模型 DEQ（即 MDEQ-Vision/）网络。由于这两类任务需要不同的输入处理方式和损失目标，它们并不直接共享训练框架。

然而，两个框架共享相同的工具代码，例如：

• lib/solvers.py：高级不动点求解器（例如 Anderson 加速和 Broyden 方法）
• lib/jacobian.py：Jacobian 相关估计（例如 Hutchinson 估计器和幂迭代法（Power method））
• lib/optimization.py：正则化方法（例如权重归一化和变分 Dropout）
• lib/layer_utils.py：层工具函数

此外，与旧版本相比，本仓库已大幅简化，便于用户进行扩展。特别地，

定理 2（“单层”DEQ 的普适性，非正式表述）：堆叠多个 DEQ
（即使使用不同类型的变换）并不会比单个 DEQ 提供更强的表示能力。

（详见论文中的正式表述。）根据上述定理，设计更好的 DEQ 模型归结为设计更优且稳定的变换函数 f_\theta。创建并尝试 DEQ 非常简单，我们推荐遵循以下 3 个步骤（本仓库也采用此流程）：

步骤 1：定义一个希望迭代至均衡状态的层 f=f_\theta

通常，这与任何深度网络层类似，应为 torch.nn.Module 的子类。该层的前向计算需要隐藏单元 z 和输入注入 x；例如：

class Layer(nn.Module):
    def __init__(self, ...):
	...
    def forward(self, z, x, **kwargs):
        return new_z

步骤 2：准备用于 DEQ 模型的不动点求解器

由于 DEQ 模型可以使用任意黑盒根求解器，我们在 lib/solvers.py 中提供了 PyTorch 实现的不动点求解器 anderson(...) 和 broyden(...)，它们会返回一个包含优化过程基本信息的字典。默认情况下，我们使用相对残差差值（即 |f(z)-z|/|z|）作为迭代停止的判据。

前向传播可简化为两行代码：

with torch.no_grad():
    # x 是输入注入；z0 是不动点的初始估计。
    z_star = self.solver(lambda z: f(z, x, *args), z0, threshold=f_thres)['result']

注意，前向传播无需存储任何中间状态，因此我们将其置于 torch.no_grad() 块中。

步骤 3：通过自动微分机制实现隐式微分

最后，我们需要确保能够计算 DEQ 的反向传播，这依赖于隐函数定理（implicit function theorem）。为此，我们可以使用 PyTorch 中的 register_hook 函数，在反向传播时注册一个后向钩子（backward hook）函数。如论文所述，反向传播本质上是求解一个涉及均衡点处 Jacobian 矩阵的线性系统的不动点：

new_z_star = self.f(z_star.requires_grad_(), x, *args)

def backward_hook(grad):
    if self.hook is not None:
        self.hook.remove()
        torch.cuda.synchronize()   # 避免无限递归
    # 求解 yJ + grad 的不动点，其中 J=J_f 是 f 在 z_star 处的 Jacobian 矩阵
    new_grad = self.solver(lambda y: autograd.grad(new_z_star, z_star, y, retain_graph=True)[0] + grad, \
                           torch.zeros_like(grad), threshold=b_thres)['result']
    return new_grad

self.hook = new_z_star.register_hook(backward_hook)

（可选）附加步骤：Jacobian 正则化（雅可比正则化）

DEQ 模型的不动点形式意味着其稳定性直接由平衡点处的 Jacobian 矩阵 J_f 所刻画。因此，我们提供了用于分析和正则化 Jacobian 性质的代码（基于 ICML'21 论文 Stabilizing Equilibrium Models by Jacobian Regularization）。具体而言，我们在训练脚本中添加了以下参数：

• jac_loss_weight: Jacobian 正则化的强度，此处我们对 ||J_f||_F（Frobenius 范数）进行正则化。
• jac_loss_freq: 随机 Jacobian 正则化的频率 p（即在训练过程中以概率 p 应用该损失）。
• jac_incremental: 若大于 0，则每经过 jac_incremental 个训练步后，将 jac_loss_weight 增加 0.1。
• spectral_radius_mode: 若为 True，则在验证集评估时估计 DEQ 模型的谱半径（spectral radius）。

因此，一个完整的 DEQ 模型实现如下所示：

from lib.solvers import anderson, broyden
from lib.jacobian import jac_loss_estimate

class DEQModel(nn.Module):
    def __init__(self, ...):
        ...
        self.f = Layer(...)
        self.solver = broyden
        ...
    
    def forward(self, x, ..., **kwargs):
        z0 = torch.zeros(...)

        # 前向传播
        with torch.no_grad():
            z_star = self.solver(lambda z: self.f(z, x, *args), z0, threshold=f_thres)['result']   # 见上文第 2 步
            new_z_star = z_star

        # （准备）反向传播，见上文第 3 步
        if self.training:
            new_z_star = self.f(z_star.requires_grad_(), x, *args)
            
            # Jacobian 相关计算，见上述附加步骤。例如：
            jac_loss = jac_loss_estimate(new_z_star, z_star, vecs=1)

            def backward_hook(grad):
                if self.hook is not None:
                    self.hook.remove()
                    torch.cuda.synchronize()   # 避免无限递归
                # 计算 yJ + grad 的不动点，其中 J=J_f 是 f 在 z_star 处的 Jacobian
                new_grad = self.solver(lambda y: autograd.grad(new_z_star, z_star, y, retain_graph=True)[0] + grad, \
                                       torch.zeros_like(grad), threshold=b_thres)['result']
                return new_grad

            self.hook = new_z_star.register_hook(backward_hook)
        return new_z_star, ...

不动点求解器（Fixed-point Solvers）

我们在 lib/solvers.py 中提供了两种通用求解器的 PyTorch 实现：broyden(...)（基于 Broyden 方法）和 anderson(...)（基于 Anderson 加速）。这两个函数接收一个变换函数 f（我们希望求解其不动点），并返回如下格式的字典：

{
 "result": ... (最接近不动点的估计值),
 "nstep": ... (给出该估计值的迭代步数),
 "abs_trace": ... (轨迹上的绝对残差),
 "rel_trace": ... (轨迹上的相对残差),
 ...
}

预训练模型

请参见 DEQ-Sequence/ 和 MDEQ-Vision/ 子目录中的链接。

致谢

• Transformer 的实现以及额外模块（例如自适应嵌入）基于 Transformer-XL 仓库。

• 本仓库的部分工具代码（例如模型摘要和 YAML 处理）修改自 HRNet 仓库。

• 我们还增加了 RAdam 优化器作为训练选项（但未设为默认）。RAdam 的实现来自 RAdam 仓库。

引用（Bibtex）

如果您在研究中使用了本仓库，请考虑引用我们的工作：

1. Deep Equilibrium Models

@inproceedings{bai2019deep,
  author    = {Shaojie Bai and J. Zico Kolter and Vladlen Koltun},
  title     = {Deep Equilibrium Models},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2019},
}

2. Multiscale Deep Equilibrium Models

@inproceedings{bai2020multiscale,
  author    = {Shaojie Bai and Vladlen Koltun and J. Zico Kolter},
  title     = {Multiscale Deep Equilibrium Models},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2020},
}

3. Stabilizing Equilibrium Models by Jacobian Regularization

@inproceedings{bai2021stabilizing,
  title     = {Stabilizing Equilibrium Models by Jacobian Regularization},
  author    = {Shaojie Bai and Vladlen Koltun and J. Zico Kolter},
  booktitle = {International Conference on Machine Learning (ICML)},
  year      = {2021}
}

DEQ（Deep Equilibrium Models）快速上手指南

环境准备

• Python 版本：≥ 3.6
• PyTorch 版本：≥ 1.10
• 硬件建议：强烈推荐使用 4 块 GPU 以获得最佳训练效率
• 网络环境：若在国内，建议配置 PyPI 镜像源（如清华源）加速依赖安装

安装步骤

1. 克隆仓库：

   git clone https://github.com/locuslab/deq.git
   cd deq
   

2. 安装依赖（推荐使用虚拟环境）：

   pip install torch>=1.10 torchvision
   # 如需处理特定任务数据（如 NLP 或视觉），请参考对应子目录中的 requirements
   

💡 国内用户可使用清华源加速安装：

pip install torch>=1.10 torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

DEQ 模型的核心思想是通过求解隐式不动点方程 $ z = f_\theta(z, x) $ 来替代传统深度网络的逐层前向传播。构建一个 DEQ 模型只需三步：

步骤 1：定义变换函数 f_\theta

创建一个继承自 torch.nn.Module 的模块，其 forward 方法接受当前状态 z 和输入 x：

import torch.nn as nn

class Layer(nn.Module):
    def __init__(self, ...):
        ...
    def forward(self, z, x, **kwargs):
        return new_z  # 返回下一个状态

步骤 2：前向传播 —— 求解不动点

使用内置求解器（如 broyden 或 anderson）在 torch.no_grad() 中求解平衡点：

from lib.solvers import broyden

# 在模型 forward 中
with torch.no_grad():
    z_star = self.solver(lambda z: self.f(z, x, *args), z0, threshold=f_thres)['result']

步骤 3：反向传播 —— 隐式微分

通过 PyTorch 的 register_hook 实现基于隐函数定理的梯度计算：

new_z_star = self.f(z_star.requires_grad_(), x, *args)

def backward_hook(grad):
    if self.hook is not None:
        self.hook.remove()
        torch.cuda.synchronize()
    new_grad = self.solver(
        lambda y: torch.autograd.grad(new_z_star, z_star, y, retain_graph=True)[0] + grad,
        torch.zeros_like(grad),
        threshold=b_thres
    )['result']
    return new_grad

self.hook = new_z_star.register_hook(backward_hook)

完整最小示例

from lib.solvers import broyden
from lib.jacobian import jac_loss_estimate
import torch
import torch.nn as nn

class DEQModel(nn.Module):
    def __init__(self, ...):
        super().__init__()
        self.f = Layer(...)
        self.solver = broyden

    def forward(self, x, f_thres=30, b_thres=30):
        z0 = torch.zeros_like(x)  # 初始猜测

        # 前向：求不动点
        with torch.no_grad():
            z_star = self.solver(lambda z: self.f(z, x), z0, threshold=f_thres)['result']

        if self.training:
            # 为反向传播准备计算图
            new_z_star = self.f(z_star.requires_grad_(), x)
            
            # 可选：Jacobian 正则化
            jac_loss = jac_loss_estimate(new_z_star, z_star, vecs=1)

            # 注册反向钩子
            def backward_hook(grad):
                if hasattr(self, 'hook') and self.hook is not None:
                    self.hook.remove()
                    torch.cuda.synchronize()
                new_grad = self.solver(
                    lambda y: torch.autograd.grad(new_z_star, z_star, y, retain_graph=True)[0] + grad,
                    torch.zeros_like(grad),
                    threshold=b_thres
                )['result']
                return new_grad

            self.hook = new_z_star.register_hook(backward_hook)
            return new_z_star, jac_loss
        else:
            return z_star

📌 提示：完整训练代码和数据预处理请参考 DEQ-Sequence/（序列任务）或 MDEQ-Vision/（视觉任务）子目录。