Tracing Computation Density in LLMs
作者: Corentin Kervadec, Iuliia Lysova, Iuri Macocco, Marco Baroni, Gemma Boleda
分类: cs.CL, cs.AI, cs.LG
发布日期: 2026-05-26
💡 一句话要点
提出s-Trace方法，揭示LLM计算密度分布规律与模块化组织结构。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 计算密度 子图搜索 模型压缩 模块化组织
📋 核心要点
- 大型语言模型计算图庞大，但并非所有计算都被有效利用，存在冗余。
- 提出s-Trace方法，通过寻找最优子图来估计LLM的有效计算密度。
- 实验表明LLM计算具有模块化结构，早期层提供粗略预测，后期层进行细化。
📝 摘要（中文）
基于Transformer的大型语言模型（LLM）由数十亿参数构成，形成深度和宽度的计算图，但尚不清楚它们是否充分利用了所有输入的能力。我们提出了一种名为s-Trace的方法，用于高效估计最能近似完整模型输出的大小为s的子图。通过这种方法，我们发现各种LLM中的计算组织成两个不同的阶段。一个主要由早期层节点组成的小子图可以重建完整模型输出分布的头部。添加更多的节点，主要位于后面的层中，并且越来越多地由注意力头组成，从而逐步改进对完整输出分布的近似。此外，我们发现每个输入所需的计算量与模型的不确定性相关，并且更稀疏的子图编码了浅层统计信息，例如一元频率。总的来说，我们的结果表明，有效的LLM计算中存在一致的模块化组织，一个稀疏的早期层核心提供了一个粗略的预测，然后通过后期层中更密集的计算进行进一步细化。
🔬 方法详解
问题定义：论文旨在解决大型语言模型（LLM）计算效率的问题。尽管LLM拥有庞大的参数量，但并非所有参数和计算都对最终输出有同等贡献。现有的方法难以有效识别和量化LLM中真正重要的计算部分，导致模型优化和压缩困难。因此，如何确定LLM中哪些计算是关键的，以及这些计算是如何组织的，是本研究要解决的核心问题。
核心思路：论文的核心思路是通过寻找一个能够最佳近似完整模型输出的子图，来估计LLM的有效计算密度。这个子图被认为是模型进行预测时所必需的关键计算部分。通过分析这个子图的结构和组成，可以了解LLM是如何利用其计算资源的，以及不同层和模块在预测过程中所扮演的角色。这种方法允许研究者追踪计算密度，并揭示LLM内部的计算组织方式。
技术框架：s-Trace方法的核心在于寻找一个大小为s的子图，该子图的输出与完整模型的输出尽可能接近。具体流程如下：1. 子图搜索：使用贪心算法或其他优化方法，从完整计算图中选择s个节点构成子图。2. 输出近似：使用选定的子图进行前向传播，得到子图的输出。3. 相似度评估：将子图的输出与完整模型的输出进行比较，使用KL散度等指标评估相似度。4. 迭代优化：重复步骤1-3，不断调整子图的节点选择，直到找到一个能够最佳近似完整模型输出的子图。
关键创新：该方法最重要的创新点在于提出了一种有效估计LLM计算密度的框架。与传统的模型分析方法不同，s-Trace方法不是简单地分析模型的参数或激活值，而是直接寻找能够最佳近似模型输出的计算子图。这种方法能够更准确地反映模型在进行预测时所实际使用的计算资源，从而更好地理解模型的内部工作机制。此外，该方法还揭示了LLM计算的模块化组织结构，为模型优化和压缩提供了新的思路。
关键设计：s-Trace方法中的关键设计包括：1. 子图大小s的选择：s的选择会影响子图的复杂度和近似精度。需要根据具体的模型和任务进行调整。2. 相似度评估指标的选择：KL散度是一种常用的评估概率分布相似度的指标，但也可以根据具体情况选择其他指标。3. 子图搜索算法的选择：贪心算法是一种简单有效的搜索算法，但也可以使用更高级的优化算法，如遗传算法或强化学习算法。
🖼️ 关键图片
📊 实验亮点
实验结果表明，LLM的计算组织成两个阶段：早期层的小型子图重建输出分布的头部，后期层通过注意力头进行细化。每个输入所需的计算量与模型的不确定性相关，稀疏子图编码浅层统计信息。这些发现揭示了LLM计算的模块化和分层特性。
🎯 应用场景
该研究成果可应用于LLM的压缩与加速，通过识别和保留关键计算子图，减少模型规模和计算量，同时保持性能。此外，该方法有助于理解LLM的内部工作机制，为模型设计和优化提供指导。未来可应用于模型蒸馏、知识迁移等领域，提升LLM在资源受限环境下的应用能力。
📄 摘要（原文）
Transformer-based large language models (LLMs) are comprised of billions of parameters arranged in deep and wide computational graphs, but it is not clear that they exploit their full capacity for all inputs. We introduce the s-Trace method to efficiently estimate the subgraph of size s that best approximates a full model output. With this method, we find the computation in a variety of LLMs to be organized in two distinct phases. A small subgraph mostly composed of early-layer nodes can reconstruct the head of the full model output distribution. Adding further nodes, mostly located in later layers and increasingly consisting of attention heads, leads to incremental refinements in approximating the full output distribution. We find moreover that the amount of necessary computation per input correlates with model uncertainty, and that sparser subgraphs encode shallow statistics, such as unigram frequency. Overall, our results suggest a consistent modular organization in effective LLM computation, with a sparse early-layer core providing a rough prediction that is further refined through denser computations in later layers.