[论文解读] Uncovering mesa-optimization algorithms in Transformers

Johannes von Oswald, Eyvind Niklasson+10 more|arXiv (Cornell University)|Sep 11, 2023

Topic Modeling被引用 5

一句话总结

本文表明，通过自回归训练的Transformer模型在推理过程中隐式地实现了基于梯度的内层优化（mesa-optimization），其中内部目标函数由上下文构建，并被优化以生成预测。作者反向工程了这一过程，提出了新型的mesa-layer，该层显式地在上下文中求解优化问题，并展示了其在少样本上下文学习和序列建模任务中性能的提升。

ABSTRACT

Transformers have become the dominant model in deep learning, but the reason for their superior performance is poorly understood. Here, we hypothesize that the strong performance of Transformers stems from an architectural bias towards mesa-optimization, a learned process running within the forward pass of a model consisting of the following two steps: (i) the construction of an internal learning objective, and (ii) its corresponding solution found through optimization. To test this hypothesis, we reverse-engineer a series of autoregressive Transformers trained on simple sequence modeling tasks, uncovering underlying gradient-based mesa-optimization algorithms driving the generation of predictions. Moreover, we show that the learned forward-pass optimization algorithm can be immediately repurposed to solve supervised few-shot tasks, suggesting that mesa-optimization might underlie the in-context learning capabilities of large language models. Finally, we propose a novel self-attention layer, the mesa-layer, that explicitly and efficiently solves optimization problems specified in context. We find that this layer can lead to improved performance in synthetic and preliminary language modeling experiments, adding weight to our hypothesis that mesa-optimization is an important operation hidden within the weights of trained Transformers.

研究动机与目标

• 探究自回归训练的Transformer是否在前向传播中隐式实现内层优化——即在未进行元训练以支持少样本学习的情况下，仍存在内部优化过程。
• 反向工程简单的序列建模任务，以揭示注意力机制如何从上下文中构建内部训练集并优化目标函数。
• 评估此类内层优化是否使标准自回归Transformer具备上下文学习能力，且不依赖于模型规模。
• 设计并评估一种新型注意力机制——mesa-layer，其显式且高效地求解上下文中的最小二乘优化问题。
• 在合成任务和初步语言建模任务中测试mesa-layer的性能，评估其提升上下文学习潜力的可行性。

提出的方法

• 对在简单序列任务上训练的自回归Transformer进行反向工程，以识别其构建训练数据和优化目标的内部机制。
• 识别早期自注意力层通过分组和复制输入标记来构建内部数据集，从而隐式定义目标函数。
• 利用深层网络执行基于梯度的优化，以生成预测。
• 提出mesa-layer，一种新型自注意力机制，通过闭式解直接求解最小二乘优化问题，而非依赖梯度步长。
• 应用键和查询的归一化以稳定mesa-layer的训练，提升收敛性和性能。
• 在GPT-2风格架构中用mesa-layer替换标准自注意力层，并在The Pile数据集上采用标准自回归语言建模范式进行评估。

实验结果

研究问题

• RQ1 自回归训练的Transformer是否在推理过程中隐式实现基于梯度的内层优化，即使未进行元学习？
• RQ2 Transformer中的内部优化过程是否可被反向工程，以揭示其如何从上下文中构建并求解内部目标？
• RQ3 标准自回归Transformer在多大程度上由于内层优化而表现出上下文中的少样本学习能力？
• RQ4 像mesa-layer这样的新型注意力机制是否能通过显式求解上下文中的优化问题，在序列建模中超越标准自注意力？
• RQ5 用mesa-layer替换标准注意力是否能提升语言建模任务中上下文学习的性能？

主要发现

• 在序列建模任务上训练的自回归Transformer实现了两步内层优化过程：(i) 通过注意力机制构建内部训练数据；(ii) 对内部目标函数进行基于梯度的优化以生成预测。
• 训练后Transformer的前向传播中包含一个隐式梯度下降算法，作用于上下文构建的目标函数，证明了标准大语言模型中存在内层优化。
• 即使是在标准自回归损失下训练的小型简单Transformer，也表现出强大的上下文少样本学习能力，且不依赖于模型规模。
• 提示微调（prompt-tuning）可提升这些模型的上下文学习性能，与大型语言模型中的作用一致。
• 所提出的mesa-layer通过直接求解最小二乘优化问题，在合成序列任务上优于标准softmax和线性自注意力Transformer。
• 初步的语言建模实验表明，用mesa-layer替换标准注意力可提升The Pile数据集上的性能，并展现出强大的上下文学习能力。