正规化流训练实践

构建正规化流涉及将可逆层及其雅可比行列式等架构组件组合成一个完整的训练流水线。这一过程包括构建 PyTorch 模型、实现精确的最大似然目标函数、通过去量化 (quantization)处理离散数据，以及通过逆向传递生成合成样本。

准备模型架构

标准的正规化流模型可以看作是一系列可逆变换的容器。它需要定义一个基分布，通常是像各向同性高斯分布这样的简单连续分布。

在训练过程中评估模型时，我们将数据输入这些层。每一层执行其变换，并返回结果以及其雅可比矩阵的对数行列式。通过累加所有层的对数行列式，我们可以计算出整个网络引起的总体体积变化。

以下是正规化流容器类的一个典型实现：

import torch
import torch.nn as nn
from torch.distributions import MultivariateNormal

class NormalizingFlow(nn.Module):
    def __init__(self, layers, num_features):
        super().__init__()
        self.layers = nn.ModuleList(layers)

        # 定义标准正态基分布
        self.register_buffer('loc', torch.zeros(num_features))
        self.register_buffer('cov', torch.eye(num_features))

    @property
    def base_distribution(self):
        return MultivariateNormal(self.loc, self.cov)

    def forward(self, x):
        # 累加雅可比矩阵的对数行列式
        log_det_jacobian = torch.zeros(x.shape[0], device=x.device)

        for layer in self.layers:
            x, ldj = layer(x)
            log_det_jacobian += ldj

        return x, log_det_jacobian

    def inverse(self, z):
        # 将样本反向通过网络以进行生成
        for layer in reversed(self.layers):
            z = layer.inverse(z)
        return z

数据去量化 (quantization)

正规化流旨在对连续概率密度函数建模。如果直接将其应用于离散数据（例如取值为 0 到 255 整数的数字图像像素），会导致模型为这些特定的整数点分配无限大的密度。这种现象会使训练过程崩溃。

为了解决这个问题，我们采用去量化。均匀去量化会向离散输入数据添加连续的均匀噪声，有效地将每个离散值的概率质量扩展到连续的单位区间上。

def uniform_dequantize(x):
    """
    向离散数据点添加连续均匀噪声。
    假设 x 包含离散整数值。
    """
    noise = torch.rand_like(x)
    return x + noise

对于更高级的应用，变分去量化使用辅助神经网络 (neural network)来学习最优噪声分布，但对于大多数初级流模型，均匀噪声已经足够有效且是标准做法。

实现损失函数 (loss function)

训练正规化流依赖于精确的最大似然估计。我们的目标是使训练数据在所学分布下的概率最大化。在实际操作中，优化算法通常是最小化某个目标，因此我们最小化负对数似然。

对于通过变换 $f$ 映射到基变量 $z$ 的数据点 $x$ ，其精确对数似然的数学表达式为：

$\log p_X(x) = \log p_Z(f(x)) + \log \left| \det \frac{\partial f(x)}{\partial x} \right|$

将此等式转化为 PyTorch 代码，需要评估基分布在变换后输出处的密度，并加上累加的雅可比对数行列式。

def compute_loss(model, x):
    # 1. 对离散输入进行均匀去量化
    x_continuous = uniform_dequantize(x)

    # 2. 将数据通过前向变换
    z, log_det_jacobian = model(x_continuous)

    # 3. 评估 z 在基分布下的对数概率
    log_prob_z = model.base_distribution.log_prob(z)

    # 4. 计算精确对数似然
    log_likelihood = log_prob_z + log_det_jacobian

    # 5. 返回整个批次的平均负对数似然
    return -log_likelihood.mean()

训练循环

架构、去量化 (quantization)和损失函数 (loss function)准备就绪后，训练循环遵循标准的 PyTorch 模式。由于流模型具有显式的参数 (parameter)化形式，我们不需要像生成对抗网络（GAN）那样同时协调多个网络。我们只需传入数据批次，计算负对数似然，并进行反向传播 (backpropagation)来更新耦合层的权重 (weight)。

import torch.optim as optim

def train_flow(model, dataloader, epochs=50, lr=1e-3):
    optimizer = optim.Adam(model.parameters(), lr=lr)
    model.train()

    for epoch in range(epochs):
        epoch_loss = 0.0

        for batch in dataloader:
            # 假设 batch 是一个元组，第一个元素是数据
            x = batch[0]

            optimizer.zero_grad()
            loss = compute_loss(model, x)
            loss.backward()

            # 建议对流模型进行梯度裁剪
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

            optimizer.step()
            epoch_loss += loss.item()

        print(f"轮次 {epoch+1}/{epochs}, 损失: {epoch_loss / len(dataloader):.4f}")

通过密度估计训练正规化流与通过逆变换生成新样本的数据处理路径对比。

采样与评估

训练循环收敛后，流变换的参数 (parameter)被优化，从而将基正态分布拉伸或压缩成训练数据的形状。

为了评估模型的生成能力，我们执行逆向传递。我们从标准正态基分布中抽取随机噪声，并将其反向推过各层。

@torch.no_grad()
def generate_samples(model, num_samples, num_features):
    model.eval()

    # 从标准正态分布中抽取随机噪声
    z = torch.randn(num_samples, num_features, device=next(model.parameters()).device)

    # 将噪声通过逆变换
    generated_data = model.inverse(z)

    return generated_data

当生成的样本与原始训练数据集的统计特性和视觉分布紧密匹配时，说明正规化流训练成功。变换已经学会将正态分布的低密度区域映射到数据空间的稀疏区域，并将高密度区域映射到集中区域。

经过高斯基分布采样并应用逆变换后，由训练好的正规化流生成的合成数据 2D 可视化图。

这部分内容有帮助吗？