[论文解读] FIT: Far-reaching Interleaved Transformers

ting chen, Lala Li|arXiv (Cornell University)|May 22, 2023

Generative Adversarial Networks and Image Synthesis被引用 5

一句话总结

FIT 引入了一种新颖的 Transformer 架构，通过分组数据标记和自适应潜在标记，交错使用局部和全局自注意力机制，实现了长序列上的高效、二次方复杂度注意力。通过将局部窗口注意力与一小部分学习得到的潜在标记上的全局注意力相结合，FIT 将计算复杂度降低至 O(L⁴/³)，同时在 16GB 显存内实现对 160K 标记输入（例如 6400×6400 图像）的端到端训练，展示了在高分辨率图像任务中无需模型并行化或优化的最先进性能。

ABSTRACT

We present FIT: a transformer-based architecture with efficient self-attention and adaptive computation. Unlike original transformers, which operate on a single sequence of data tokens, we divide the data tokens into groups, with each group being a shorter sequence of tokens. We employ two types of transformer layers: local layers operate on data tokens within each group, while global layers operate on a smaller set of introduced latent tokens. These layers, comprising the same set of self-attention and feed-forward layers as standard transformers, are interleaved, and cross-attention is used to facilitate information exchange between data and latent tokens within the same group. The attention complexity is $O(n^2)$ locally within each group of size $n$, but can reach $O(L^{{4}/{3}})$ globally for sequence length of $L$. The efficiency can be further enhanced by relying more on global layers that perform adaptive computation using a smaller set of latent tokens. FIT is a versatile architecture and can function as an encoder, diffusion decoder, or autoregressive decoder. We provide initial evidence demonstrating its effectiveness in high-resolution image understanding and generation tasks. Notably, FIT exhibits potential in performing end-to-end training on gigabit-scale data, such as 6400$ imes$6400 images, or 160K tokens (after patch tokenization), within a memory capacity of 16GB, without requiring specific optimizations or model parallelism.

研究动机与目标

• 解决标准 Transformer 在处理长序列时的二次方计算复杂度问题。
• 在不牺牲性能且无需模型并行化的情况下，实现在长序列上的高效、可扩展注意力机制。
• 设计一种统一架构，可作为编码器、扩散解码器或自回归解码器，适用于多种视觉任务。
• 探索交错局部与全局 Transformer 层并结合自适应潜在标记在改进长距离依赖建模方面的有效性。

提出的方法

• 将输入数据划分为标记组，每组在窗口化上下文中通过局部自注意力进行处理。
• 为每组引入少量可学习的潜在标记以表示全局上下文，从而实现复杂度降低的全局自注意力。
• 交错使用局部和全局 Transformer 层，通过交叉注意力实现在数据标记与潜在标记之间的双向信息流动。
• 在单次前向传播中迭代更新数据标记与潜在标记，以确保局部与全局表征的深度整合。
• 在局部与全局层中保持标准的前馈网络与注意力机制，仅在输入范围与注意力模式上存在差异。
• 利用交叉注意力在每组内实现从数据标记到潜在标记以及反向的信息路由，实现自适应计算。

实验结果

研究问题

• RQ1 结合潜在标记的交错局部与全局 Transformer 层是否能比标准或稀疏注意力机制实现更优的长距离建模？
• RQ2 每组中潜在标记的数量如何影响高分辨率图像任务中的模型性能与效率？
• RQ3 与顺序堆叠或仅使用局部层相比，交错局部与全局层是否能提升性能？
• RQ4 FIT 是否能在不使用模型并行化或特殊优化的情况下，于 16GB 显存内实现对千兆级输入（如 6400×6400 图像）的端到端训练？
• RQ5 FIT 在视觉领域中作为自回归、扩散与编码器任务的通用架构，其适用范围在多大程度上具有普适性？

主要发现

• FIT 在 153M 参数下实现了 ImageNet-64×64 自回归建模的接近最先进性能，达到 3.42 比特/维度，与更大模型相当。
• 增加潜在标记数量可提升性能（降低负对数似然与比特每维），且对参数量与训练速度影响极小，表明其具有高效的可扩展性。
• 交错局部与全局层相比顺序堆叠或仅使用局部层，能实现更优性能（如 ImageNet 上为 3.49 比特每维），且训练成本几乎无增加。
• FIT 实现了在 16GB 显存内对 160K 标记序列（如 6400×6400 图像）的端到端训练，无需模型并行化或专用优化。
• 该架构在任务间表现出强大的泛化能力，在高分辨率视觉基准中可有效作为编码器、扩散解码器与自回归解码器使用。