[论文解读] Object Representations as Fixed Points: Training Iterative Refinement Algorithms with Implicit Differentiation

Michael Chang, Thomas L. Griffiths+1 more|arXiv (Cornell University)|Jul 2, 2022

Domain Adaptation and Few-Shot Learning被引用 5

一句话总结

本文提出使用隐式微分训练迭代精炼模型（如槽注意力机制），以学习以对象为中心的表征，通过稳定的、可微分的不动点公式替代反向传播通过展开的迭代过程。该方法实现了更优的优化稳定性、更低的验证损失和更优的生成质量，且内存和时间复杂度保持恒定，仅需额外一行代码。

ABSTRACT

Iterative refinement -- start with a random guess, then iteratively improve the guess -- is a useful paradigm for representation learning because it offers a way to break symmetries among equally plausible explanations for the data. This property enables the application of such methods to infer representations of sets of entities, such as objects in physical scenes, structurally resembling clustering algorithms in latent space. However, most prior works differentiate through the unrolled refinement process, which can make optimization challenging. We observe that such methods can be made differentiable by means of the implicit function theorem, and develop an implicit differentiation approach that improves the stability and tractability of training by decoupling the forward and backward passes. This connection enables us to apply advances in optimizing implicit layers to not only improve the optimization of the slot attention module in SLATE, a state-of-the-art method for learning entity representations, but do so with constant space and time complexity in backpropagation and only one additional line of code.

研究动机与目标

• 解决在潜在集合表征学习的迭代精炼模型中，通过展开迭代进行反向传播所导致的不稳定性和高计算成本问题。
• 利用槽注意力等迭代精炼算法的不动点特性，实现隐式微分，从而解耦前向传播与反向传播过程。
• 在不修改网络架构的前提下，提升最先进模型（如 SLATE 和原始槽注意力）的训练效率与稳定性。
• 证明隐式微分可实现鲁棒训练，无需梯度裁剪、学习率预热或对迭代次数进行超参数调优。
• 展示该方法在不同架构上的泛化能力，保持分割质量并提升下游对象属性预测性能。

提出的方法

• 将槽注意力模块建模为不动点迭代，其中最终表征是 z = f_x(z) 的解，从而在不动点处实现隐式微分。
• 应用隐函数定理计算梯度，无需展开迭代过程，利用一阶 Neumann 近似实现高效的反向传播。
• 将隐式梯度集成至损失计算中，支持端到端训练，相比标准反向传播仅需额外一行代码。
• 将前向传播（展开迭代）与反向传播（隐式梯度计算）解耦，使每次反向传播的内存和时间复杂度保持恒定。
• 通过使用精炼函数 f 的雅可比矩阵，利用 Neumann 级数近似计算隐式梯度，避免因展开反向传播导致的梯度爆炸。
• 将该方法应用于 SLATE 和原始槽注意力架构，验证其在不同解码器和训练设置下的泛化能力。

实验结果

研究问题

• RQ1 隐式微分能否稳定训练因展开反向传播导致梯度爆炸的迭代精炼模型（如槽注意力）？
• RQ2 隐式微分是否可在无需梯度裁剪或学习率预热的情况下，实现对以对象为中心表征学习的更好优化？
• RQ3 在使用隐式微分与标准反向传播时，精炼迭代次数对训练稳定性和性能的影响如何？
• RQ4 隐式微分能否在不同网络架构下保持生成掩码和分割输出的质量？
• RQ5 与标准槽注意力相比，隐式微分在下游任务（如对象属性预测）上的性能提升程度如何？

主要发现

• 隐式槽注意力在三个基准数据集上的训练过程中，验证损失显著低于原始槽注意力和 SLATE。
• 与 SLATE 相比，该方法将 Fréchet Inception Distance (FID) 降低最多 30%，图像重建的均方误差 (MSE) 降低最多 50%。
• 隐式槽注意力消除了对梯度裁剪、学习率预热或迭代次数超参数调优的需求，而这些在标准训练中是必需的。
• 该方法在反向传播中保持恒定的空间和时间复杂度，与精炼迭代次数无关。
• 在 CLEVR 数据集上，隐式槽注意力使用七次迭代时，重建 MSE 比使用一次迭代的原始槽注意力低 2 倍。
• 隐式槽注意力在不同架构中均保持高质量的分割掩码，包括 Locatello 等人 [39] 提出的原始空间广播解码器，如图 8 所示。