归一化流与变分推断

概述

归一化流（Normalizing Flows）是一种通过可逆变换构造复杂概率分布的方法。¹ 其核心思想是将简单的基础分布（如高斯分布）通过一系列可逆映射变换为复杂的目标分布。

在变分推断的框架下，归一化流提供了更灵活的后验近似族，使得我们可以用更复杂的分布来逼近真实后验，从而提高变分推断的精度。

---

变量变换公式

定理

设 $\mathbf{z}\in {\mathbb{R}}^d$ 服从分布 $p(\mathbf{z})$，$\mathbf{f}:{\mathbb{R}}^d\to {\mathbb{R}}^d$ 是双射（可逆且光滑），令 $\mathbf{x}=\mathbf{f}(\mathbf{z})$，则：

$$p_X(\mathbf{x})=p_Z({\mathbf{f}}^{-1}(\mathbf{x}))\cdot \left|\det \frac{\partial {\mathbf{f}}^{-1}}{\partial \mathbf{x}}\right|$$

或等价地：

$$\log p_X(\mathbf{x})=\log p_Z(\mathbf{z})+\log \left|\det \frac{\partial {\mathbf{f}}^{-1}}{\partial \mathbf{x}}\right|,\mathbf{z}={\mathbf{f}}^{-1}(\mathbf{x})$$

对数行列式的性质

$$\log \left|\det \frac{\partial {\mathbf{f}}^{-1}}{\partial \mathbf{x}}\right|=-\log \left|\det \frac{\partial \mathbf{f}}{\partial \mathbf{z}}\right|$$

因此，我们更常用：

$$\log p_X(\mathbf{x})=\log p_Z({\mathbf{f}}^{-1}(\mathbf{x}))-\log \left|\det \frac{\partial \mathbf{f}}{\partial \mathbf{z}}\right|$$

雅可比行列式计算的核心挑战

计算 $\det (\partial \mathbf{f}/\partial \mathbf{z})$ 的复杂度为 $O(d^3)$，其中 $d$ 是维度。因此，设计高效的归一化流需要精心设计变换，使雅可比行列式易于计算。

---

归一化流的基本构建块

1. 仿射变换（Affine Transformation）

最简单的变换形式：

$$\mathbf{f}(\mathbf{z})=\mathbf{a}+\mathbf{b}\odot \mathbf{z}$$

其中 $\mathbf{a}$ 是平移向量，$\mathbf{b}$ 是缩放向量。

特点：

• 雅可比行列式：${\prod }_d{b}_d$
• 计算复杂度：$O(d)$
• 表达能力有限

2. 逐通道仿射耦合（Affine Coupling）

由Dinh et al. (2015)在Real NVP中提出。²

结构：

$$\begin{array}{rl}{\mathbf{x}}_{1:d}& ={\mathbf{z}}_{1:d}\\ {\mathbf{x}}_{d+1:D}& ={\mathbf{z}}_{d+1:D}\odot \exp (s({\mathbf{z}}_{1:d}))+t({\mathbf{z}}_{1:d})\end{array}$$

其中 $s(\cdot )$ 和 $t(\cdot )$ 是任意神经网络。

雅可比行列式：

$$\det \frac{\partial \mathbf{f}}{\partial \mathbf{z}}=\prod _{i=d+1}^D\exp (s_i({\mathbf{z}}_{1:d}))=\exp \left(\sum _{i=d+1}^D{s}_i({\mathbf{z}}_{1:d})\right)$$

计算复杂度：$O(1)$（对角矩阵）

3. 置换层（Permutation）

通过固定置换操作来混合变量：

class Permutation(nn.Module):
    def __init__(self, dim, permutation=None):
        super().__init__()
        if permutation is None:
            self.permutation = torch.randperm(dim)
        else:
            self.permutation = permutation
    
    def forward(self, z):
        return z[:, self.permutation]
    
    def inverse(self, x):
        inv_permutation = torch.argsort(self.permutation)
        return x[:, inv_permutation]

4. 激活归一化（ActNorm）

通过数据依赖的初始化实现高效训练：

$$\mathbf{x}=\mathbf{z}\odot \mathbf{w}+\mathbf{b}$$

其中 $\mathbf{w},\mathbf{b}$ 通过数据批量初始化，使输出分布均值为0、方差为1。

---

主流归一化流架构

Real NVP (Real-valued Non-Volume Preserving)

由Dinh et al. (2015)提出，是首个成功的深度归一化流。²

架构：

class RealNVP(nn.Module):
    def __init__(self, dim, hidden_dim=512):
        super().__init__()
        self.mask = self._create_mask(dim)
        # 两层网络用于s和t
        self.scale_net = nn.Sequential(
            nn.Linear(dim // 2, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, dim // 2)
        )
        self.translate_net = nn.Sequential(
            nn.Linear(dim // 2, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, dim // 2)
        )
    
    def _create_mask(self, dim):
        mask = torch.zeros(dim)
        mask[:dim // 2] = 1
        return mask
    
    def forward(self, z):
        # 分割输入
        z1 = z * self.mask
        z2 = z * (1 - self.mask)
        # 计算仿射参数
        s = self.scale_net(z1)
        t = self.translate_net(z1)
        # 应用变换
        x2 = z2 * torch.exp(s) + t
        # 合并
        x = self.mask * x + (1 - self.mask) * x2
        # 返回雅可比行列式
        log_det = torch.sum(s * (1 - self.mask))
        return x, log_det

特点：

• 可并行计算
• 易于存储和反转
• 需要多层堆叠以增加表达能力

Glow (Generative Flow with Invertible 1x1 Convolutions)

由Kingma & Dhariwal (2018)提出，改进了Real NVP。³

关键创新：可逆的1x1卷积

class Invertible1x1Conv(nn.Module):
    def __init__(self, dim):
        super().__init__()
        # 初始化为随机正交矩阵
        self.W = nn.Parameter(torch.linalg.qr(torch.randn(dim, dim))[0])
    
    def forward(self, z):
        log_det = torch.logdet(self.W) * z.shape[1]  # 批次维度
        return torch.einsum('bij,jk->bik', z, self.W), log_det
    
    def inverse(self, x):
        W_inv = torch.inverse(self.W)
        return torch.einsum('bij,jk->bik', x, W_inv)

Masked Autoregressive Flow (MAF)

基于自回归模型构建归一化流。⁴

变换：

$$x_i=z_i\cdot \exp (s_i({\mathbf{z}}_{1:i-1}))+t_i({\mathbf{z}}_{1:i-1})$$

特点：

• 雅可比为三角矩阵，对角元素为 $\exp (s_i({\mathbf{z}}_{1:i-1}))$
• 采样容易（逐样本生成）
• 密度估计需要顺序计算（慢）

Neural Spline Flows (NSF)

使用分段线性/三次样条作为非线性变换。⁵

B-样条变换：

$$f(x)=\sum _{k=0}^K{w}_k{B}_k(x)$$

其中 $B_k$ 是B-样条基函数。

优势：

• 表达能力更强
• 可以精确逆变换（通过查表）
• 在密度估计任务上效果优异

---

归一化流与变分推断

增强变分后验

标准VAE使用高斯后验近似：

$$q_{\varphi }(\mathbf{z}\mid \mathbf{x})=\mathcal{N}(\mathbf{z};{\mu }_{\varphi }(\mathbf{x}),{\sigma }_{\varphi }^2(\mathbf{x}))$$

使用归一化流可以增强后验近似：

$$\mathbf{z}={\mathbf{f}}_K(\cdots {\mathbf{f}}_1({\mathbf{z}}_0)\cdots ),{\mathbf{z}}_0\sim \mathcal{N}({\mu }_{\varphi }(\mathbf{x}),{\sigma }_{\varphi }^2(\mathbf{x}))$$

class FlowVAE(nn.Module):
    def __init__(self, encoder, decoder, flow_layers):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.flow = PlanarFlow(num_layers=flow_layers)
    
    def forward(self, x):
        # 编码
        q0_mean, q0_logvar = self.encoder(x)
        q0_std = torch.exp(0.5 * q0_logvar)
        
        # 从高斯采样
        z0 = q0_mean + q0_std * torch.randn_like(q0_mean)
        
        # 通过流变换
        zK, log_det = self.flow(z0)
        
        # ELBO
        log_pzK = self.prior.log_prob(zK)
        log_qz0 = -0.5 * (q0_logvar + (z0 - q0_mean)**2 / q0_std**2)
        log_px_z = self.decoder(zK)
        
        elbo = log_px_z + log_pzK - log_qz0 + log_det
        return elbo.mean()

Planar流

最简单的可学习流变换：

$$\mathbf{f}(\mathbf{z})=\mathbf{z}+\mathbf{u}\cdot h({\mathbf{w}}^T\mathbf{z}+b)$$

其中 $h$ 是激活函数（如tanh）。

径向流（Radial Flow）

另一种常用变换：

$$\mathbf{f}(\mathbf{z})=\mathbf{z}+\frac{\beta }{\alpha +\Vert \mathbf{z}-{\mathbf{z}}_0\Vert }(\mathbf{z}-{\mathbf{z}}_0)$$

ELBO中的流变换

使用归一化流后，ELBO变为：

$$\mathcal{L}(\theta ,\varphi )={\mathbb{E}}_{q_{\varphi }({\mathbf{z}}_0\mid \mathbf{x})}\left[\log p_{\theta }(\mathbf{x}\mid \mathbf{f}({\mathbf{z}}_0))\right]+{\mathbb{E}}_{q_{\varphi }({\mathbf{z}}_0\mid \mathbf{x})}\left[\log p({\mathbf{z}}_0)-\log q_{\varphi }({\mathbf{z}}_0\mid \mathbf{x})+\log \left|\det \frac{\partial \mathbf{f}}{\partial {\mathbf{z}}_0}\right|\right]$$

---

连续归一化流（CNF）

神经微分方程视角

将离散的流层推广为连续的变换过程：

$$\frac{d\mathbf{z}(t)}{dt}=f_{\theta }(\mathbf{z}(t),t)$$

概率流ODE

$$\frac{d\mathbf{x}}{dt}=f_{\theta }(\mathbf{x},t)$$

可以证明，这对应于某个确定性随机过程的轨迹。

FFJORD (Free-form Jacobian of Reversible Dynamics)

使用Jacobian-vector积避免显式计算雅可比行列式。⁶

class FFJORD(nn.Module):
    def __init__(self, dim, hidden_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim + 1, hidden_dim),  # +1 for time
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, dim)
        )
    
    def drift(self, z, t):
        # f(z, t) 用于 dz/dt = f(z, t)
        return self.net(torch.cat([z, t * torch.ones_like(z[:, :1])], dim=-1))
    
    def forward(self, z0, t_span):
        # 使用ODE求解器
        solution = odeint(self.drift, z0, t_span)
        return solution[-1]

---

实现框架

normflows库

import normflows as nf
 
# 定义基础分布
base = nf.distributions.base.DiagGaussian(2)
 
# 定义流变换
flows = [
    nf.flows.AffineCouplingBlock(
        nf.nets.MLP([1, 64, 64, 1])
    ),
    nf.flows.Permute(2),
    # ... 更多层
]
 
# 构建流模型
nf_model = nf.NormalizingFlow(base=base, flows=flows)
 
# 训练
optimizer = torch.optim.Adam(nf_model.parameters(), lr=1e-4)
for data in dataloader:
    optimizer.zero_grad()
    loss = -nf_model.log_prob(data).mean()
    loss.backward()
    optimizer.step()

nflows库

另一个流行的归一化流实现库。

---

归一化流的优缺点

优点

优点	说明
精确对数密度计算	避免变分近似的误差
可逆性	易于从隐空间采样和重建
可解释性	变换路径提供了分布转化的直观理解
灵活的表达能力	通过堆叠多层增加表达能力

缺点

缺点	说明
高计算成本	需要$O(d^3)$或精心设计的变换
表达能力限制	仍然受限于可逆变换的结构
大内存开销	保存所有中间激活用于逆向传播

---

应用场景

1. 变分自编码器增强

使用归一化流增强VAE的后验近似，提高生成质量。

2. 图像生成

Glow在高质量图像生成上的应用。

3. 密度估计

NSF在表格数据和低维数据上的密度估计。

4. 数据增强

利用流的逆变换生成高质量合成样本。

---

参考

---

相关主题

• variational-inference-advanced - 变分推断进阶
• probabilistic-circuits-fundamentals - 概率电路基础
• diffusion-variational - 扩散模型变分推断
• generative-models-survey - 生成模型综述
• bayesian-neural-networks-advanced-inference - 贝叶斯神经网络高级推断

Footnotes

1. Kobyzev et al. (2021). “Normalizing Flows: An Introduction and Review of Current Methods”. IEEE TPAMI. ↩

2. Dinh et al. (2015). “Density Estimation Using Real NVP”. ICLR 2017. ↩ ↩²

3. Kingma & Dhariwal (2018). “Glow: Generative Flow with Invertible 1x1 Convolutions”. NeurIPS 2018. ↩

4. Papamakarios et al. (2017). “Masked Autoregressive Flow for Density Estimation”. NeurIPS 2017. ↩

5. Durkan et al. (2019). “Neural Spline Flows”. NeurIPS 2019. ↩

6. Grathwohl et al. (2019). “FFJORD: Free-form Jacobian of Reversible Dynamics”. ICML 2019. ↩