iFlame: Interleaving Full and Linear Attention for Efficient Mesh Generation

📄 arXiv: 2503.16653v2 📥 PDF

作者: Hanxiao Wang, Biao Zhang, Weize Quan, Dong-Ming Yan, Peter Wonka

分类: cs.CV

发布日期: 2025-03-20 (更新: 2025-03-24)

备注: Project website: https://wanghanxiao123.github.io/iFa/

💡 一句话要点

iFlame：交错式全注意力和线性注意力的高效网格生成方法

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 网格生成 Transformer 注意力机制 线性注意力 3D建模

📋 核心要点

  1. 现有基于注意力的网格生成模型计算复杂度高，难以处理高分辨率3D数据，而线性注意力模型又难以捕捉长程依赖。
  2. iFlame采用交错式框架，结合线性注意力的效率和全注意力的表达能力，在沙漏架构中实现高效网格生成。
  3. 实验表明，iFlame在ShapeNet和Objaverse数据集上，能高效生成高质量3D网格，并显著提升推理速度。

📝 摘要（中文）

本文提出了一种名为iFlame的、基于Transformer的新型网格生成网络架构。虽然基于注意力的模型在网格生成方面表现出了卓越的性能，但其二次计算复杂度限制了可扩展性，尤其是在处理高分辨率3D数据时。相反，线性注意力机制虽然计算成本较低，但通常难以捕捉长程依赖关系，导致次优的结果。为了解决这种权衡，我们提出了一种交错自回归网格生成框架，该框架结合了线性注意力的效率和全注意力机制的表达能力。为了进一步提高效率并利用网格表示的固有结构，我们将这种交错方法集成到沙漏架构中，从而显著提高了效率。我们的方法减少了训练时间，同时实现了与纯粹基于注意力的模型相当的性能。为了提高推理效率，我们实现了一种缓存算法，与原始Transformer相比，该算法几乎使速度提高了一倍，并将KV缓存大小减少了七分之八。我们在ShapeNet和Objaverse上评估了我们的框架，证明了其高效生成高质量3D网格的能力。我们的结果表明，所提出的交错框架有效地平衡了计算效率和生成性能，使其成为网格生成的一种实用解决方案。在Objaverse上，使用4个GPU，在包含最多4k面的39k数据上，训练仅需2天。

🔬 方法详解

问题定义：论文旨在解决3D网格生成任务中，基于Transformer的注意力机制计算复杂度高，难以扩展到高分辨率网格的问题。现有方法要么计算量大，要么无法有效捕捉长程依赖关系，导致生成质量下降。

核心思路：论文的核心思路是结合全注意力机制和线性注意力机制的优点，通过交错使用这两种注意力机制，在保证生成质量的同时，降低计算复杂度。同时，利用网格的固有结构，设计沙漏架构，进一步提升效率。

技术框架：iFlame采用自回归的网格生成框架，整体架构为沙漏型。编码器部分逐渐减小特征图尺寸，提取全局特征；解码器部分逐渐增大特征图尺寸，生成网格。在编码器和解码器的每一层中，交替使用全注意力机制和线性注意力机制。此外，为了加速推理过程，还设计了缓存算法，减少KV缓存的大小。

关键创新：论文的关键创新在于提出了交错式注意力机制，将全注意力和线性注意力结合使用，在计算效率和模型表达能力之间取得了平衡。此外，针对网格数据的特点，设计了沙漏型架构，进一步提升了效率。缓存算法也显著提升了推理速度。

关键设计：iFlame的关键设计包括：1) 交错式注意力机制的具体实现方式，例如全注意力和线性注意力的比例、位置等；2) 沙漏型架构的具体层数、特征图尺寸变化等；3) 缓存算法的具体实现细节，例如缓存大小、更新策略等。论文中使用了交叉熵损失函数来训练模型。具体参数设置未知。

📊 实验亮点

iFlame在ShapeNet和Objaverse数据集上进行了评估，实验结果表明，iFlame能够在保证生成质量的前提下，显著降低计算成本。在Objaverse数据集上，使用4个GPU，在包含最多4k面的39k数据上，训练仅需2天。此外，iFlame的缓存算法能够将推理速度提高近一倍，并将KV缓存大小减少七分之八。

🎯 应用场景

iFlame在3D内容生成领域具有广泛的应用前景，例如游戏开发、虚拟现实、计算机辅助设计等。它可以用于生成各种类型的3D模型，例如人物、物体、场景等。该研究的实际价值在于降低了3D模型生成的计算成本，提高了生成效率，使得高质量3D内容的生成更加容易。未来，iFlame可以进一步扩展到其他3D生成任务，例如纹理生成、材质生成等。

📄 摘要（原文）

This paper propose iFlame, a novel transformer-based network architecture for mesh generation. While attention-based models have demonstrated remarkable performance in mesh generation, their quadratic computational complexity limits scalability, particularly for high-resolution 3D data. Conversely, linear attention mechanisms offer lower computational costs but often struggle to capture long-range dependencies, resulting in suboptimal outcomes. To address this trade-off, we propose an interleaving autoregressive mesh generation framework that combines the efficiency of linear attention with the expressive power of full attention mechanisms. To further enhance efficiency and leverage the inherent structure of mesh representations, we integrate this interleaving approach into an hourglass architecture, which significantly boosts efficiency. Our approach reduces training time while achieving performance comparable to pure attention-based models. To improve inference efficiency, we implemented a caching algorithm that almost doubles the speed and reduces the KV cache size by seven-eighths compared to the original Transformer. We evaluate our framework on ShapeNet and Objaverse, demonstrating its ability to generate high-quality 3D meshes efficiently. Our results indicate that the proposed interleaving framework effectively balances computational efficiency and generative performance, making it a practical solution for mesh generation. The training takes only 2 days with 4 GPUs on 39k data with a maximum of 4k faces on Objaverse.