[论文解读] Scalable Adaptive Computation for Iterative Generation

Allan Jabri, David J. Fleet+1 more|arXiv (Cornell University)|Dec 22, 2022

Stochastic Gradient Optimization Techniques被引用 13

一句话总结

本文提出循环接口网络（RINs），一种与领域无关的架构，通过使用潜在标记来聚焦全局自注意力，将计算与输入数据维度解耦，从而实现对高维数据的自适应、可扩展生成。RINs 在图像和视频扩散建模中达到最先进性能——无需级联或引导即可扩展至 1024×1024 分辨率，同时通过潜在自条件化技术利用循环机制实现更深层次、上下文更丰富的路由，相比 2D/3D U-Net 将 FLOPs 降低高达 10 倍。

ABSTRACT

Natural data is redundant yet predominant architectures tile computation uniformly across their input and output space. We propose the Recurrent Interface Networks (RINs), an attention-based architecture that decouples its core computation from the dimensionality of the data, enabling adaptive computation for more scalable generation of high-dimensional data. RINs focus the bulk of computation (i.e. global self-attention) on a set of latent tokens, using cross-attention to read and write (i.e. route) information between latent and data tokens. Stacking RIN blocks allows bottom-up (data to latent) and top-down (latent to data) feedback, leading to deeper and more expressive routing. While this routing introduces challenges, this is less problematic in recurrent computation settings where the task (and routing problem) changes gradually, such as iterative generation with diffusion models. We show how to leverage recurrence by conditioning the latent tokens at each forward pass of the reverse diffusion process with those from prior computation, i.e. latent self-conditioning. RINs yield state-of-the-art pixel diffusion models for image and video generation, scaling to 1024X1024 images without cascades or guidance, while being domain-agnostic and up to 10X more efficient than 2D and 3D U-Nets.

研究动机与目标

• 解决深度学习模型在高维、冗余自然数据（如图像和视频）上计算均匀性低效的问题。
• 开发一种可扩展架构，根据数据内容动态分配计算，避免输入尺寸带来的二次方复杂度。
• 通过利用循环机制，在迭代生成任务中实现深层、富有表现力的路由，实现跨步骤的上下文传播 amortization（摊销）。
• 通过设计更高效、与领域无关的架构，取代像素级扩散模型中对级联或引导的依赖。

提出的方法

• 将隐藏单元分解为接口标记（随输入尺寸线性增长）和紧凑潜在标记（固定尺寸），实现计算与输入布局的解耦。
• 使用交叉注意力在接口标记与潜在标记之间传递信息，实现双向自底向上和自顶向下的反馈。
• 仅在少量潜在标记上应用全局自注意力，避免输入尺寸带来的二次方复杂度。
• 引入潜在自条件化：重用先前推理步骤中的潜在激活作为后续步骤的上下文，实现类似循环的效果，但无需时间反向传播。
• 堆叠 RIN 模块构建更深的网络，通过迭代累积上下文，提升路由的表达能力。
• 将模型以标准可微分网络形式进行训练和部署，计算图固定，支持现代硬件上的高效推理。

实验结果

研究问题

• RQ1 我们能否设计一种神经架构，在不引入输入尺寸二次方复杂度的前提下，对高维数据实现自适应计算分配？
• RQ2 在迭代生成任务中，如何利用循环机制在不进行时间反向传播的情况下，实现更深层次的路由？
• RQ3 像 RINs 这类与领域无关的架构，是否能在高分辨率图像和视频生成任务中，同时在样本质量和 FLOP 效率上超越基于 U-Net 的扩散模型？
• RQ4 潜在自条件化在迭代扩散设置中，能在多大程度上提升路由深度和模型表达能力？
• RQ5 是否可能通过更高效的架构，使像素级扩散模型在无需级联或引导的情况下，实现 1024×1024 分辨率的扩展？

主要发现

• RINs 在 64×64 到 1024×1024 各分辨率下均在 ImageNet 上达到最先进 FID 分数，且未使用引导或级联机制。
• 在 1024×1024 图像生成任务中，RINs 相较于 2D 和 3D U-Net 基线，每推理步骤的 FLOPs 最多降低 10 倍。
• 在 Kinetics600 视频生成基准上，RINs 的表现优于领先方法，同时每步 FLOPs 降低 10 倍。
• 潜在自条件化通过在迭代去噪步骤中保留路由上下文，实现了更深的有效计算图，提升了模型表达能力。
• RINs 在多种模态上保持强大性能，具有领域无关性，仅依赖注意力和 MLP，具有广泛适用性。
• 该模型仅通过极少的架构修改即实现优异性能，表明自适应计算可高效集成至标准扩散训练流程中。