[论文解读] Progress measures for grokking via mechanistic interpretability

Neel Nanda, Lawrence Chan+3 more|arXiv (Cornell University)|Jan 12, 2023

Neural Networks and Applications被引用 54

一句话总结

这篇论文对解决模加法的小型变换器进行反向工程，以揭示基于傅里叶变换的算法，并定义进展度量（受限损失和排除损失），以表明 grokking 源自逐步形成电路后再通过权重衰减消除记忆。

ABSTRACT

Neural networks often exhibit emergent behavior, where qualitatively new capabilities arise from scaling up the amount of parameters, training data, or training steps. One approach to understanding emergence is to find continuous extit{progress measures} that underlie the seemingly discontinuous qualitative changes. We argue that progress measures can be found via mechanistic interpretability: reverse-engineering learned behaviors into their individual components. As a case study, we investigate the recently-discovered phenomenon of ``grokking'' exhibited by small transformers trained on modular addition tasks. We fully reverse engineer the algorithm learned by these networks, which uses discrete Fourier transforms and trigonometric identities to convert addition to rotation about a circle. We confirm the algorithm by analyzing the activations and weights and by performing ablations in Fourier space. Based on this understanding, we define progress measures that allow us to study the dynamics of training and split training into three continuous phases: memorization, circuit formation, and cleanup. Our results show that grokking, rather than being a sudden shift, arises from the gradual amplification of structured mechanisms encoded in the weights, followed by the later removal of memorizing components.

研究动机与目标

• 激发寻找平滑、具有因果联系的进展度量，以揭示神经网络中出现的行为背后的原因。
• 证明 grokking 可以通过对学习到的电路进行机制性反向工程来解释。
• 对模加法 Transformer 进行反向工程，以揭示傅里叶乘法算法。
• 定义并验证连续的进展度量，为 grokking 相变铺垫。
• 将训练动态表征为三阶段：记忆化、电路形成和清理。

提出的方法

• 用 P=113 的模 P 加法对小型 Transformer 模型进行训练。
• 反向工程权重和激活，确定一个基于傅里叶的加法算法。
• 显示嵌入、注意力和 MLP 激活在一组关键频率处呈现周期性结构。
• 证明解嵌出层和 MLP 层实现三角恒等式以计算 a+b mod P。
• 将 restricted loss 和 excluded loss 定义为进展度量并进行实证验证。

实验结果

研究问题

• RQ1 小型 Transformer 训练在模加法上的 grokking 背后存在哪些机制结构？
• RQ2 我们能否识别一个连续、可解释的进展度量，先于 grokking 转变？
• RQ3 学习到的算法在频率分量和三角恒等式方面如何运作？
• RQ4 从记忆到泛化的训练动态（阶段）有哪些特征？
• RQ5 权重衰减在推动 grokking 和相变中的作用是什么？

主要发现

• 模型将输入嵌入到一个圆上，使用傅里叶分量通过三角恒等式执行加法。
• 对数映射 W_L 可以很好地近似为五个关键频率之和，从而实现傅里叶乘法读出。
• 大多数神经元可以被单一频率的二次多项式很好近似，且对数值的读出是频率局部化的。
• 消融实验表明关键频率是必要的；移除非关键频率可以提高性能。
• grokk ing 包含三个阶段——记忆化、电路形成和清理，权重衰减引导向泛化电路的转变。
• 提出的进展度量（restricted loss 与 excluded loss）在 grokking 之前持续增加，为训练动态提供了透视。