[论文解读] Memory in Plain Sight: Surveying the Uncanny Resemblances of Associative Memories and Diffusion Models

Benjamin Hoover, Hendrik Strobelt+4 more|arXiv (Cornell University)|Sep 28, 2023

Neural Networks and ApplicationsComputer Science被引用 3

一句话总结

本文揭示了扩散模型与联想记忆之间深层的数学等价性，表明两者均作为梯度下降过程，在李雅普诺夫能量函数上收敛至不动点吸引子。通过将扩散模型框架化为类似于霍普菲尔德网络的能量系统，该研究为收敛性、记忆容量和优化提供了理论保证，同时在统一的动力系统框架下整合了生成建模与记忆检索。

ABSTRACT

The generative process of Diffusion Models (DMs) has recently set state-of-the-art on many AI generation benchmarks. Though the generative process is traditionally understood as an "iterative denoiser", there is no universally accepted language to describe it. We introduce a novel perspective to describe DMs using the mathematical language of memory retrieval from the field of energy-based Associative Memories (AMs), making efforts to keep our presentation approachable to newcomers to both of these fields. Unifying these two fields provides insight that DMs can be seen as a particular kind of AM where Lyapunov stability guarantees are bypassed by intelligently engineering the dynamics (i.e., the noise and step size schedules) of the denoising process. Finally, we present a growing body of evidence that records DMs exhibiting empirical behavior we would expect from AMs, and conclude by discussing research opportunities that are revealed by understanding DMs as a form of energy-based memory.

研究动机与目标

• 通过动力系统与常微分方程（ODEs），建立扩散模型与联想记忆之间的正式数学联系。
• 证明扩散模型的反向去噪过程等价于联想记忆中的能量梯度下降。
• 揭示联想记忆所具备的理论优势（如保证收敛至不动点），而这些在标准扩散模型中缺失。
• 通过联想记忆理论的视角重新诠释扩散模型与Transformer架构，激发生成建模领域的新研究方向。
• 通过能量系统中的记忆容量，为理解大模型中的缩放定律提供理论基础。

提出的方法

• 将扩散模型形式化为连续时间ODE，揭示其与李雅普诺夫能量函数梯度下降的等价性。
• 将扩散模型的反向去噪过程映射至联想记忆的记忆检索动力学，后者收敛至能量景观的局部极小值。
• 利用基于得分模型推导出的能量函数定义李雅普诺夫函数，确保反向过程中稳定性和收敛性。
• 将深度学习中成熟的优化技术（如ADAM和L-BFGS）应用于扩散模型的采样过程，将其视为能量上的梯度下降。
• 将类比扩展至Transformer模型，表明注意力机制类似于稠密联想记忆的单步更新，进而提出“能量Transformer”模块。
• 从能量景观中存储的不动点数量角度，分析记忆容量与缩放定律，将模型规模与性能与联想记忆容量联系起来。

Figure 1: Comparing the emphases of Diffusion Models and Associative Memories tasked with learning the same energy (negative log-probability) landscape, represented with both contours and gradient arrows. Diffusion Models (left) train a score function (depicted as orange arrows) to model the gradien

实验结果

研究问题

• RQ1 扩散模型的反向去噪动力学与联想记忆的记忆检索动力学之间有何关系？
• RQ2 联想记忆中固有的理论保证（如收敛性）在标准扩散模型中为何缺失？
• RQ3 扩散模型的采样过程能否被重新诠释为能量函数上的梯度下降？其优化优势为何？
• RQ4 现代架构（如Transformer）在计算动力学上在多大程度上类似于联想记忆？
• RQ5 如何通过联想记忆理论的视角，理论化理解大模型中的记忆容量与缩放定律？

主要发现

• 扩散模型的反向去噪过程在数学上等价于李雅普诺夫能量函数上的梯度下降，与联想记忆完全一致。
• 与扩散模型不同，联想记忆保证收敛至不动点（局部极小值），提供了理论上的稳定性优势。
• 联想记忆中的能量函数可直接计算与优化，从而可引入标准扩散模型中不可用的新分析与训练技术。
• 扩散模型可被重新诠释为收敛至吸引子的连续时间动力系统，从而可理论表征记忆容量与稳定性。
• 通过将采样过程视为能量梯度下降，可利用ADAM和L-BFGS等先进优化方法加速扩散模型的采样过程。
• Transformer的注意力机制与稠密联想记忆的单步更新高度相似，暗示注意力与记忆检索之间存在更深层次的理论联系。

Figure 2: Diffusion Models through the lens of PF-ODEs, adapted from [ 2 ] . The forward (corrupting) process takes a true data point at time $s=0$ and corrupts the data into pure noise at time $s=T$ using an SDE. The reverse (reconstructing) process undoes the corruption using the score of the dist