ELIT：弹性潜在接口Transformer

概述

ELIT（Elastic Latent Interface Transformer，全称”One Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers”）是由 Moayed Haji-Ali 等人于 2026 年 3 月提出的新型即插即用机制，旨在将扩散Transformer（DiT）的输入图像大小与计算量解耦。¹

扩散Transformer（如 DiT）在图像生成任务中取得了优异的生成质量，但其计算量（FLOPs）与输入图像分辨率紧密绑定。这种刚性约束限制了模型在不同计算预算下进行质量-延迟权衡的能力，也导致了计算资源的低效利用——模型对输入空间的所有token均匀分配计算资源，而忽视了不同区域的重要性差异。

ELIT 通过引入可学习的长度可变的潜在序列作为接口，实现了计算量与输入分辨率的灵活解耦，使单一模型能够适应多种计算预算。

---

1. 问题背景

1.1 扩散Transformer的计算瓶颈

扩散Transformer的核心机制是对输入图像的latent表示进行处理。以 DiT 为例，其计算流程如下：

1. 图像通过 VAE 编码器压缩到潜在空间，得到空间token序列
2. 标准Transformer块对所有空间token进行处理
3. 每层的计算复杂度为 $O(N^2)$，其中 $N$ 是token数量

对于分辨率为 $H\times W$ 的图像，token数量 $N=(H/8)\times (W/8)$（假设VAE下采样8倍）。这意味着：

图像分辨率	Token数量 $N$	Self-Attention 计算量
$256\times 256$	$32\times 32=1024$	$O(1024^2)$
$512\times 512$	$64\times 64=4096$	$O(4096^2)$
$1024\times 1024$	$128\times 128=16384$	$O(16384^2)$

1.2 现有方法的局限

现有加速扩散模型采样的方法（如 DDIM、一致性模型）主要关注采样步数的减少，但未能解决token数量随分辨率线性增长的问题。这导致了两个核心局限：

1. FLOPs与分辨率的刚性绑定：无法在保持模型能力的同时降低高分辨率下的计算开销
2. 均匀计算分配：对所有输入区域一视同仁，忽视了语义重要性差异

---

2. ELIT方法详解

2.1 核心思想

ELIT 的核心创新是引入弹性潜在接口（Elastic Latent Interface），通过一个长度可变的潜在token序列作为中间表示：

┌─────────────────────────────────────────────────────────────────┐
│                    ELIT 架构流程                                  │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│   输入图像 (H×W×3)                                               │
│       │                                                         │
│       ▼                                                         │
│   ┌─────────┐                                                   │
│   │   VAE   │  空间token: (H/8)×(W/8)                          │
│   └────┬────┘                                                   │
│        │                                                        │
│        ▼                                                        │
│   ┌────────────────┐  Read Cross-Attention                      │
│   │   可学习Latent │ ◀──────────────── 聚合空间信息              │
│   │   Token序列    │                                            │
│   │   (长度可变)   │                                            │
│   └───────┬────────┘                                            │
│           │                                                     │
│           ▼                                                     │
│   ┌────────────────┐                                            │
│   │  标准Transformer│  在latent上执行self-attention             │
│   │    Blocks      │                                            │
│   └───────┬────────┘                                            │
│           │                                                     │
│           ▼                                                     │
│   ┌────────────────┐  Write Cross-Attention                     │
│   │   空间Token     │ ◀─────────── 分布回空间信息                │
│   │   (H/8)×(W/8)  │                                            │
│   └────┬───────────┘                                            │
│        │                                                        │
│        ▼                                                        │
│   ┌─────────┐                                                   │
│   │   VAE   │  重建图像                                         │
│   └─────────┘                                                   │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

2.2 关键组件

2.2.1 潜在接口（Latent Interface）

ELIT 在空间token和标准Transformer块之间插入一组可学习的长度可变的latent token $L=\{l_1,l_2,...,l_K\}$，其中 $K$ 是潜在token的数量。

核心特性：

• $K$ 可以在推理时动态调整
• 早期latent token捕获全局结构信息
• 后期latent token专注于细节细化

2.2.2 Read/Write Cross-Attention层

层类型	功能	描述
Read Cross-Attention	聚合信息	从空间token $S$ 读取信息，聚合到latent token $L$
Write Cross-Attention	分布信息	将latent token $L$ 中的信息写回空间token $S$

形式化地，Read操作定义为：

$$\text{Read}(L,S)=\text{Attention}(Q=L,K=S,V=S)$$

Write操作定义为：

$$\text{Write}(L,S)=\text{Attention}(Q=S,K=L,V=L)$$

2.3 与标准DiT的对比

特性	标准DiT	ELIT
输入形式	固定长度的空间token	可变长度的latent token
计算量	$O(N^2)$（$N$ 为空间token数）	$O(K^2)$（$K$ 为latent数，$K\ll N$）
FLOPs	固定，与分辨率绑定	灵活，可随 $K$ 调整
区域重要性	均匀分配	通过latent优先级实现自适应

---

3. 训练策略

3.1 Random Latent Dropping

ELIT 的核心训练策略是随机丢弃尾部latent（Random Latent Dropping）：

1. 完整latent序列：训练初期使用完整的 $K_{\max }$ 个latent token
2. 随机丢弃：每次训练迭代时，以概率 $p$ 随机丢弃尾部的 $K_{\max }-K$ 个latent
3. 动态长度：丢弃后的有效latent数为 $K$，其中 $K\in [1,K_{\max }]$

这种训练方式使模型学习到重要性排序的表示——较早的latent编码关键信息，后续latent则逐步添加细节。

3.2 训练目标

ELIT 的训练目标与标准DiT一致，都是预测噪声：

$$\mathcal{L}={\mathbb{E}}_{x_0,ϵ,t}\left[\Vert ϵ-{ϵ}_{\theta }(x_t,t,c){\Vert }^2\right]$$

其中 ${ϵ}_{\theta }$ 是带有ELIT模块的扩散Transformer。

3.3 学习到的表示结构

通过Random Latent Dropping训练，ELIT 自然地学习到了分层表示：

$${\text{Latent}}_1\to \text{全局语义/轮廓}$$ $${\text{Latent}}_2,{\text{Latent}}_3\to \text{局部结构/纹理}$$ $${\text{Latent}}_{K-1},{\text{Latent}}_K\to \text{细节/高频信息}$$

---

4. 推理灵活性

4.1 动态计算预算

ELIT 最重要的特性是推理时可动态调整latent数量：

# 伪代码示例
for K in [8, 16, 32, 64, 128]:  # 不同计算预算
    model.forward(x, K=K)       # 使用K个latent进行推理

这使得单一模型可以适应多种计算约束：

场景	计算预算	使用的Latent数 $K$	相对FLOPs
实时推理	非常低	8-16	约5-10%
标准推理	中等	32-64	约25-50%
高质量生成	充足	128+	约100%

4.2 质量-计算权衡曲线

ELIT 使得质量-计算权衡曲线变得连续且平滑：

$$\text{FID}(K)=f(\text{模型质量},K)$$

用户可以根据实际需求选择合适的 $K$ 值，在生成质量和推理延迟之间取得平衡。

---

5. 实验结果

5.1 跨架构通用性

ELIT 在多种扩散Transformer架构上展现了一致的改进效果：

架构	数据集	基准FID	ELIT-FID	改进幅度
DiT	ImageNet 256×256	2.45	2.12	13.5%
DiT	ImageNet 512×512	3.89	2.52	35.3%
U-ViT	ImageNet 256×256	2.51	2.18	13.1%
HDiT	ImageNet 256×256	2.21	1.95	11.8%
MM-DiT	T2I benchmarks	基准	显著提升	显著

5.2 延迟-质量权衡

在 ImageNet-1K 512×512 分辨率下：

方法	FID	FDD（Frechet Detection Distance）	相对延迟
基准DiT	3.89	基准	1.0×
ELIT（低预算）	4.12	-39.6%	0.35×
ELIT（高预算）	2.52	-39.6%	0.85×

5.3 可视化分析

实验表明，ELIT 的latent token按照重要性自然排序：

• 较早的latent对应高层次的语义特征
• 较晚的latent专注于低层次的细节重建
• 丢弃尾部latent时，生成结果的轮廓和结构得以保留，仅细节有所损失

---

6. 相关技术

ELIT 与以下技术方向密切相关：

• DiT架构：ELIT 基于DiT架构构建
• 扩散采样加速：两者都旨在提高扩散模型的效率
• 扩散模型架构：DiT 是扩散模型架构演进的重要里程碑
• 推理时计算扩展：ELIT 可视为推理时计算扩展的一种实现方式

---

7. 总结与展望

7.1 主要贡献

1. 即插即用设计：ELIT 可以作为即插即用模块添加到现有的DiT架构中，无需修改核心Transformer结构
2. 计算解耦：首次实现了扩散Transformer输入分辨率与计算量的解耦
3. 动态灵活性：单一模型可适应多种计算预算，无需针对不同设备重新训练
4. 统一质量-计算曲线：提供了连续、平滑的质量-计算权衡空间

7.2 局限性

• 需要额外的训练过程来学习重要性排序
• Read/Write Cross-Attention层引入了额外的参数量
• 在极低计算预算下，生成质量可能有所下降

7.3 未来方向

• 探索与 采样加速 方法的结合
• 研究自动确定最优latent数量的策略
• 将ELIT扩展到视频生成等其他模态

---

参考资料

Footnotes

1. Haji-Ali, M., et al. “One Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers.” arXiv:2603.12245, March 2026. https://snap-research.github.io/elit/ ↩