Weight-based Decomposition: A Case for Bilinear MLPs
作者: Michael T. Pearce, Thomas Dooms, Alice Rigg
分类: cs.LG, cs.AI
发布日期: 2024-06-06
💡 一句话要点
基于权重的分解：双线性MLP的案例，提升模型可解释性
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 模型可解释性 双线性层 张量分解 特征提取 深度学习
📋 核心要点
- 现代模型可解释性差，难以理解特征与模型权重之间的关系，阻碍了模型改进和调试。
- 论文提出一种基于权重的分解方法，将双线性层分解为稀疏交互的特征向量，以增强模型的可解释性。
- 实验表明，该方法在浅层图像分类和小型语言模型中展现出良好的可解释性，并可将现有模型微调为双线性变体。
📝 摘要（中文）
门控线性单元（GLU）已成为现代基础模型中的常见构建块。双线性层去除了“门”中的非线性，但仍具有与其他GLU相当的性能。双线性层的一个吸引人的特性是，它们可以完全用三阶张量和线性运算来表示。利用这一点，我们开发了一种方法，将双线性张量分解为一组稀疏交互的特征向量，这些特征向量在浅层图像分类器（MNIST）和小型语言模型（Tiny Stories）的初步实验中显示出良好的可解释性。由于分解完全等同于模型的原始计算，因此双线性层可能是一种对可解释性友好的架构，有助于将特征与模型权重联系起来。我们的方法的应用可能不限于预训练的双线性模型，因为我们发现诸如TinyLlama-1.1B之类的语言模型可以被微调为双线性变体。
🔬 方法详解
问题定义：现有深度学习模型，特别是基于门控线性单元（GLU）的模型，虽然性能强大，但其内部运作机制复杂，难以理解。特征与模型权重之间的关系模糊，导致模型的可解释性较差。这使得模型的调试、改进以及在安全敏感领域的应用面临挑战。双线性层虽然简化了GLU的结构，但其可解释性仍有待提升。
核心思路：论文的核心思路是将双线性层表示为三阶张量，并对其进行分解，提取出稀疏交互的特征向量。通过分析这些特征向量，可以更好地理解模型内部的特征表示和计算过程，从而提高模型的可解释性。这种分解方法旨在将复杂的模型权重转化为更易于理解的特征表示。
技术框架：该方法首先将双线性层转化为等价的三阶张量表示。然后，利用张量分解技术，将该张量分解为一组稀疏交互的特征向量。这些特征向量代表了模型学习到的主要特征模式。最后，通过分析这些特征向量，可以理解模型如何利用这些特征进行预测。整个框架保持了与原始模型的等价性，确保分解后的模型能够复现原始模型的计算结果。
关键创新：该方法最重要的创新点在于将双线性层与张量分解技术相结合，提出了一种新的模型可解释性分析方法。与传统的模型可解释性方法不同，该方法直接从模型权重出发，通过分解提取出具有物理意义的特征向量，从而更直接地揭示模型内部的运作机制。此外，该方法还表明，现有的语言模型可以通过微调转化为双线性变体，从而为该方法在更广泛的模型中的应用提供了可能性。
关键设计：论文的关键设计包括：1) 将双线性层表示为三阶张量的具体数学形式；2) 选择合适的张量分解算法，以提取稀疏交互的特征向量；3) 设计实验验证分解后的特征向量的可解释性，例如通过可视化或与领域知识进行对比；4) 研究如何将现有模型微调为双线性变体，并评估其性能。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在MNIST图像分类和Tiny Stories语言模型中展现出良好的可解释性。通过分解双线性层，可以提取出具有物理意义的特征向量，例如在MNIST中可以识别数字的笔画，在Tiny Stories中可以识别语法结构。此外，实验还证明了可以将TinyLlama-1.1B等语言模型微调为双线性变体，而性能没有显著下降。
🎯 应用场景
该研究成果可应用于提升深度学习模型的可解释性，尤其是在图像分类和自然语言处理领域。通过理解模型内部的特征表示，可以更好地调试和改进模型，并提高模型在安全敏感领域的应用可靠性。此外，该方法还可以用于知识发现，从模型权重中提取有用的领域知识。
📄 摘要（原文）
Gated Linear Units (GLUs) have become a common building block in modern foundation models. Bilinear layers drop the non-linearity in the "gate" but still have comparable performance to other GLUs. An attractive quality of bilinear layers is that they can be fully expressed in terms of a third-order tensor and linear operations. Leveraging this, we develop a method to decompose the bilinear tensor into a set of sparsely interacting eigenvectors that show promising interpretability properties in preliminary experiments for shallow image classifiers (MNIST) and small language models (Tiny Stories). Since the decomposition is fully equivalent to the model's original computations, bilinear layers may be an interpretability-friendly architecture that helps connect features to the model weights. Application of our method may not be limited to pretrained bilinear models since we find that language models such as TinyLlama-1.1B can be finetuned into bilinear variants.