MAR: 一种无需向量量化的自回归图像生成模型

MAR模型简介

MAR(Masked Autoregressive)是由MIT和Meta AI联合提出的一种新型自回归图像生成模型。与传统的自回归模型不同,MAR摒弃了向量量化(VQ)的使用,直接在连续的VAE潜在空间中进行自回归建模,从而能够生成更高质量、更多样化的图像。

MAR生成的图像示例

MAR模型的核心思想是将图像生成过程分解为两个阶段:

使用预训练的VAE将图像编码到低维潜在空间
在潜在空间中使用自回归模型逐步生成图像表示

这种方法避免了传统自回归模型中VQ带来的信息损失,能够更好地捕捉图像的细节和结构。

MAR的主要特点

1. 无需向量量化

MAR最大的创新点在于摒弃了向量量化。传统自回归模型如VQ-VAE需要将连续的潜在表示离散化,这inevitably会导致信息损失。MAR直接在连续空间中建模,保留了更多细节信息。

2. 高效的训练和推理

MAR采用了一种称为DiffLoss的新型损失函数,能够在训练过程中更有效地学习像素间的依赖关系。在推理阶段,MAR使用了一种新的采样策略,大大提高了生成速度。

3. 强大的生成能力

在ImageNet 256x256数据集上的实验表明,MAR在FID和Inception Score等指标上都达到了新的state-of-the-art水平。特别是MAR-H模型,在50K样本上达到了1.55的FID分数,展现了出色的图像生成质量。

4. 灵活的架构设计

MAR提供了多个不同规模的模型版本(MAR-B、MAR-L、MAR-H),可以根据具体应用场景和计算资源选择合适的模型。同时,MAR的架构也支持进一步的扩展和优化。

MAR的实现细节

MAR的PyTorch实现主要包含以下几个关键组件:

VAE编码器: 使用预训练的VAE将输入图像编码到16维的潜在空间。
自回归Transformer: 核心的自回归模型,用于在潜在空间中逐步生成图像表示。
DiffLoss: 新设计的损失函数,用于更有效地学习像素间的依赖关系。
采样策略: 包括classifier-free guidance等技术,用于提高生成质量和多样性。

以下是MAR模型的核心代码片段:

class MAR(nn.Module):
    def __init__(self, ...):
        super().__init__()
        self.vae = ...  # 预训练的VAE
        self.transformer = ...  # 自回归Transformer
        self.diffloss = DiffLoss(...)  # DiffLoss模块
    
    def forward(self, x):
        z = self.vae.encode(x)  # VAE编码
        pred = self.transformer(z)  # 自回归预测
        loss = self.diffloss(pred, z)  # 计算DiffLoss
        return loss

MAR的训练与评估

MAR的训练过程采用了分布式数据并行(DDP)技术,以充分利用多GPU资源。以MAR-L模型为例,其默认训练配置如下:

训练轮数: 400 epochs
学习率: 1e-4 (使用warmup)
BatchSize: 64 (每个GPU)
DiffLoss配置: 3个MLP块,宽度1024

训练脚本示例:

torchrun --nproc_per_node=8 --nnodes=4 \
main_mar.py \
--model mar_large --diffloss_d 3 --diffloss_w 1024 \
--epochs 400 --warmup_epochs 100 --batch_size 64 --blr 1.0e-4 \
--output_dir ${OUTPUT_DIR} --data_path ${IMAGENET_PATH}