[论文解读] The Dual Form of Neural Networks Revisited: Connecting Test Time Predictions to Training Patterns via Spotlights of Attention

Kazuki Irie, Róbert Csordás+1 more|arXiv (Cornell University)|Feb 11, 2022

Adversarial Robustness in Machine Learning被引用 5

一句话总结

本文重新审视神经网络中线性层的对偶形式，将其表达为对所有训练数据点和初始权重的未归一化点积注意力机制。通过在测试时计算注意力权重，该方法可直接可视化影响预测的训练模式，揭示了视觉和自然语言处理任务中输入与输出之间可解释的、上下文敏感的关联，尽管在规模上存在计算限制。

ABSTRACT

Linear layers in neural networks (NNs) trained by gradient descent can be expressed as a key-value memory system which stores all training datapoints and the initial weights, and produces outputs using unnormalised dot attention over the entire training experience. While this has been technically known since the 1960s, no prior work has effectively studied the operations of NNs in such a form, presumably due to prohibitive time and space complexities and impractical model sizes, all of them growing linearly with the number of training patterns which may get very large. However, this dual formulation offers a possibility of directly visualising how an NN makes use of training patterns at test time, by examining the corresponding attention weights. We conduct experiments on small scale supervised image classification tasks in single-task, multi-task, and continual learning settings, as well as language modelling, and discuss potentials and limits of this view for better understanding and interpreting how NNs exploit training patterns. Our code is public.

研究动机与目标

• 将训练好的神经网络线性层重新解释为对整个训练数据集的注意力机制，从而实现对测试时预测结果与训练模式之间关系的直接追溯。
• 探究在此对偶形式中，注意力权重是否能揭示测试输入与特定训练样本之间的有意义且可由人类理解的关系。
• 评估该方法在小规模图像分类和语言建模范式中的可行性与可解释性。
• 研究注意力模式如何反映测试查询与训练数据之间的语义和结构相似性，特别是在少样本或分布外泛化场景下。
• 评估该方法在内存、计算和对已训练模型适用性方面的局限性。

提出的方法

• 将深度神经网络中线性层的前向传播表示为对所有训练输入-输出对的未归一化点积注意力机制，将这些对视为键（keys）和值（values）。
• 使用测试输入作为查询（query），在所有训练模式上计算注意力得分，其中注意力权重表示相关性。
• 在训练过程中记录线性层的输入，并在推理阶段计算注意力权重，以建立测试预测与训练数据之间的关联。
• 将该方法应用于MNIST和Fashion-MNIST上的前馈网络，以及WikiText-2和Aesop’s Fables等文本数据集上的LSTM模型。
• 通过可视化注意力最高的训练标记或图像块，评估注意力是否与语义或结构相似性对齐。
• 采用键-值机制，其中键为训练输入，值为训练目标，查询由测试输入生成。

实验结果

研究问题

• RQ1 在神经网络线性层的对偶形式中，注意力权重能否揭示出对测试预测贡献最大的特定训练样本？
• RQ2 被重点关注的训练模式是否与测试输入在语义或结构上相似，表明其具备上下文推理能力？
• RQ3 该对偶形式注意力机制在不同学习场景（如单任务、多任务和持续学习）下的行为表现如何？
• RQ4 该方法能否通过识别源训练样本，解释大模型中的泛化现象，如零样本翻译或少样本推理？
• RQ5 该方法在内存、计算和对真实世界模型适用性方面的实际限制是什么？

主要发现

• 在图像分类任务中，模型在测试时的注意力始终集中于具有相似类别标签和视觉模式的训练样本，即使在多任务和持续学习设置中也是如此。
• 在语言建模范式中，注意力最高的训练段落与测试查询在语义和结构上高度相似，例如在查询铁甲舰相关问题时，注意力集中在与战舰相关的文本上。
• 尽管句法结构相似但语义无关的负样本受到较低注意力，表明注意力机制具备上下文感知能力，而不仅仅是依赖句法。
• 该方法成功实现了对整个训练经验的注意力可视化，揭示了视觉和语言任务中测试输入与训练数据之间可解释的关联。
• 尽管内存和计算成本较高，该方法在小规模数据集上仍表现出可行性，预示着随着硬件进步，未来在更大模型中具有应用潜力。
• 该分析不适用于已训练好的模型，因为它需要在训练过程中记录线性层的输入，从而限制了对已有模型的回溯分析。