NNsight and NDIF: Democratizing Access to Open-Weight Foundation Model Internals

📄 arXiv: 2407.14561v4 📥 PDF

作者: Jaden Fiotto-Kaufman, Alexander R. Loftus, Eric Todd, Jannik Brinkmann, Koyena Pal, Dmitrii Troitskii, Michael Ripa, Adam Belfki, Can Rager, Caden Juang, Aaron Mueller, Samuel Marks, Arnab Sen Sharma, Francesca Lucchetti, Nikhil Prakash, Carla Brodley, Arjun Guha, Jonathan Bell, Byron C. Wallace, David Bau

分类: cs.LG, cs.AI

发布日期: 2024-07-18 (更新: 2025-04-01)

备注: Code at https://nnsight.net

💡 一句话要点

NNsight与NDIF： democratize 开放权重基础模型内部机制的研究

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 神经网络 可解释性 深度学习 远程执行 干预实验 开源工具 大型语言模型

📋 核心要点

  1. 现有研究大型神经网络内部机制的成本高昂且复杂，阻碍了相关研究的广泛开展。
  2. NNsight和NDIF通过延迟远程执行和共享GPU资源，实现了对大型模型内部机制的透明和高效访问。
  3. 实验表明，该框架能够支持多种研究方法，并提供与现有方法相比更优的性能。

📝 摘要（中文）

本文介绍了NNsight和NDIF，它们协同工作，旨在促进对超大型神经网络所学习到的表征和计算的科学研究。NNsight是一个开源系统，它扩展了PyTorch，引入了延迟远程执行。国家深度推理结构（NDIF）是一个可扩展的推理服务，它执行NNsight请求，允许用户共享GPU资源和预训练模型。这些技术由干预图（Intervention Graph）架构支持，该架构旨在将实验设计与模型运行时解耦。总而言之，该框架提供了对深度神经网络（如超大型语言模型LLM）内部机制的透明且高效的访问，而无需承担单独托管定制模型的成本或复杂性。我们对机器学习文献进行了定量调查，揭示了对大规模AI内部机制研究日益增长的差距。我们通过在大型模型上启用一系列研究方法，展示了我们框架的设计和使用，以解决这一差距。最后，我们进行了基准测试，以将性能与以前的方法进行比较。

🔬 方法详解

问题定义：现有方法研究大型神经网络内部机制时，面临着模型部署成本高、资源需求大、实验设计与模型运行耦合等问题。这些问题限制了研究人员对大型模型内部运作机制的深入探索，阻碍了相关领域的进展。

核心思路：本文的核心思路是解耦实验设计与模型运行，通过NNsight提供灵活的实验设计接口，并通过NDIF提供可扩展的远程执行服务。用户无需自行部署和运行大型模型，即可进行各种干预实验，从而降低研究成本和复杂度。

技术框架：该框架包含两个主要组件：NNsight和NDIF。NNsight是一个基于PyTorch的开源库，它引入了延迟远程执行的概念，允许用户在本地设计实验，并在远程服务器上执行。NDIF是一个可扩展的推理服务，它负责接收NNsight的请求，并在GPU集群上运行模型。干预图（Intervention Graph）是连接NNsight和NDIF的关键架构，它定义了实验的执行流程和数据传输方式。

关键创新：最重要的技术创新点在于干预图（Intervention Graph）的设计，它将实验设计与模型运行解耦，使得用户可以灵活地定义各种干预实验，而无需修改模型代码。此外，延迟远程执行也降低了本地计算资源的需求，使得更多研究人员可以参与到大型模型的内部机制研究中。

关键设计：NNsight的关键设计在于其灵活的API，允许用户方便地插入干预操作，例如修改神经元的激活值、替换模型的权重等。NDIF的关键设计在于其可扩展的架构，可以根据需求动态调整GPU资源，保证服务的稳定性和性能。具体的参数设置、损失函数、网络结构等技术细节取决于所研究的具体模型。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文通过定量分析机器学习文献，揭示了对大型AI模型内部机制研究的不足。同时，通过实验证明，NNsight和NDIF能够有效降低研究成本和复杂度，并支持多种研究方法。基准测试表明，该框架在性能上优于现有方法，为大规模AI模型的研究提供了新的工具和平台。

🎯 应用场景

该研究成果可广泛应用于大型语言模型的可解释性研究、模型调试与优化、以及安全性评估等方面。通过NNsight和NDIF，研究人员可以更方便地理解模型的内部运作机制，发现潜在的漏洞和偏差，并进行针对性的改进，从而提升模型的性能和可靠性。

📄 摘要（原文）

We introduce NNsight and NDIF, technologies that work in tandem to enable scientific study of the representations and computations learned by very large neural networks. NNsight is an open-source system that extends PyTorch to introduce deferred remote execution. The National Deep Inference Fabric (NDIF) is a scalable inference service that executes NNsight requests, allowing users to share GPU resources and pretrained models. These technologies are enabled by the Intervention Graph, an architecture developed to decouple experimental design from model runtime. Together, this framework provides transparent and efficient access to the internals of deep neural networks such as very large language models (LLMs) without imposing the cost or complexity of hosting customized models individually. We conduct a quantitative survey of the machine learning literature that reveals a growing gap in the study of the internals of large-scale AI. We demonstrate the design and use of our framework to address this gap by enabling a range of research methods on huge models. Finally, we conduct benchmarks to compare performance with previous approaches. Code, documentation, and tutorials are available at https://nnsight.net/.