小米大模型团队 2025-04-18 17:00 北京
最近,GPT-4o 生成的吉卜力风格的图像深受大家喜爱,大家在对图像的效果感到惊艳的同时,也纷纷好奇它背后使用了什么样的图像生成技术。有一种猜测认为,GPT-4o 很可能使用的是逐行扫描(raster scan order)的自回归图像生成方式,即一行一行地依次生成图像。
小米大模型团队最新提出了一种新的生成方法,即具有方向感知的对角蛇形(diagonal snake-like order)自回归图像生成方式(DAR)。简单来说,它不是一行一行地画,而是沿着图像的对角线、像蛇一样灵活地生成每个像素。这样的方式比传统方法更自然,也更接近人类绘画时的直觉。我们在 256×256 的ImageNet 基准测试中,DAR 取得了 1.37 的 FID 分数,刷新了当前同类技术的最好成绩(SoTA)。
相关的论文、训练代码、模型权重已经开源,欢迎大家交流讨论。
训练代码:https://github.com/xiaomi-research/dar
模型参数:https://huggingface.co/wbh123/dar
论文arxiv地址:https://arxiv.org/abs/2503.11129
01
为什么传统的逐行扫描方式存在局限?
自回归图像生成技术一直面临逐行扫描带来的挑战:在行尾换行时,模型需要处理较远距离的预测和生成方向的剧烈变化,导致模型难以准确高效地生成高质量的图像。如下图(a)所示,我们以换行区块5->6为例说明这种方式存在的问题。
预测距离远
输入是区块5,却要预测区块6,它们的位置编码在二维空间上距离很远。
生成方向突变
如果我们将输出和输入区块位置编码的差值定义为生成方向,不难发现换行区块的统一生成方向和预测方向是 (1, -4),其它区块的预测方向是 (0, 1)。transformer 没有显式地获得生成方向的信息,难以处理有 2 种生成方向的序列。而且,生成方向不同会导致 hidden states 的相对位置一直变化,使得 RoPE 不准确。
例如,transformer 的不同层里,h4 (4->5) 的位置在(0,3) -> (0,4) 之间变化,h5 (5->6) 的位置在 (0,4) -> (1,0) 之间变化。所有层使用 RoPE 编码 (0,3) 和 (0,4) 的相对位置是不准确的。
为解决这一难题,我们提出了一种创新性的技术方案——对角蛇形扫描顺序,如图(b)所示。相比逐行扫描,它保证了所有输入和输出区块之间的距离都尽可能小。针对生成方向会动态变化的问题,我们在注意力计算中增加了一个显性的生成方向编码作为位置编码的补充,以提升注意力的生成方向感知能力。下面是详细介绍。
02
整体架构
整个模型是decoder-only的结构,保持了和LLM兼容的next token prediction的训练和推理方法,decoder的网络结构跟LlamaGen一致。我们选择IBQ工作中的image tokenizer的codebook作为图像token embedding。我们创新地提出了4D-RoPE以及Direction Embedding来控制图像的生成方式。
给第个点设置 4 维坐标
,RoPE 的旋转复数矩阵
的计算方式为
简单来说,1D-RoPE 是用 维表示
,2D-RoPE 是分别用
维表示
,4D-RoPE 是分别用
维表示
4D-RoPE 能同时编码 2 个 token 之间的当前位置和预测位置的相对位置信息,相当于直接给注意力矩阵添加了生成方向的信息。
对角蛇形扫描顺序有 4 种生成方向:,分别对应朝下,朝右,朝右上,朝左下四个方向。用 4 个可学习的嵌入分别表示这 4 种生成方向,将 direction embedding 和 class embedding 相加计算 AdaLN 中的参数。
03
实验结果
如下图所示,我们在 256x256 ImageNet 类别条件生成任务上FID取得了SoTA的结果1.37。相比之前的SoTA方法RAR,在更小的尺寸上(1.1B vs 1.5B),也取得了更好的效果 1.42 < 1.48。尤其是随着model size的增大,生成的图像质量也越来越高。
在此基础上,我们做了如下的消融实验,我们发现方向感知模块以及codebook的使用,对逐行扫描的顺序也有明显提升。4D-RoPE以及Direction Embedding单独使用也有效,如下图所示。
在推理效率上,和一般自回归方法相同,主要受参数量以及step数影响,如下图所示。
04
总结与展望
我们提出了一种具有方向感知的对角蛇形扫描自回归图像生成框架(DAR),有效地确保相邻索引的 token 在空间上紧密相邻。此外,方向感知模块显著增强了模型在处理频繁变化的生成方向上的表现。我们开发了一系列不同规模的模型,参数量从485M到2.0B不等。在实验中,我们的DAR模型在256×256的ImageNet基准测试中取得了突破性的FID分数(1.37),超越了此前所有自回归方法。
未来,我们将进一步支持更灵活的多种分辨率图像生成。鉴于本方法与LLM的训练和推理方式高度兼容,我们还将持续探索更加统一的多模态理解与生成技术方案。我们的论文、训练代码与模型权重均已开源,期待与社区广泛交流和共同发展。