💬 NLP

TraceNAS: Zero-shot LLM Pruning via Gradient Trace Correlation

TraceNAS 是一个高效的无需训练的神经架构搜索框架，它利用梯度迹相关性来识别最优的非均匀剪枝大语言模型，通过将其损失平面与原始预训练模型对齐，以极低的计算成本实现了与训练感知方法相媲敌的性能。

原作者： Prajna G. Malettira, Manish Nagaraj, Arjun Roy, Shubham Negi, Kaushik Roy

发布于 2026-02-04

📖 1 分钟阅读☕ 轻松阅读

原作者： Prajna G. Malettira, Manish Nagaraj, Arjun Roy, Shubham Negi, Kaushik Roy

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你拥有一座庞大、极其聪明的图书馆（即大语言模型，简称 LLM），它包含了人类知识的总和。由于它过于庞大且沉重，以至于无法随身携带，也无法在普通的手机上运行。你想把它缩小成一个可以装进兜里的版本，同时又不丢失它讲故事、解谜或理解笑话的能力。

这就是 TraceNAS 这篇论文试图解决的问题。

问题所在：剪错了东西

通常，当人们试图缩小这些巨型模型时，他们的行为就像一个笨拙的园丁。他们可能会说：“让我们从每棵树上剪掉 50% 的树枝，”或者“让我们去掉最重的叶子。”这被称为均匀剪枝（uniform pruning）。

但问题在于：大脑（或图书馆）的所有部分并不都是平等的。有些部分是承载复杂逻辑的“重要器官”，而另一些部分仅仅是容易修剪掉的“脂肪”。如果你剪掉了错误的器官，模型就会丧失记忆。如果你只剪掉脂肪，它依然会过于沉重。

现有的方法试图通过以下两种方式来确定该剪掉什么：

一次只看一部分： 它们检查某个特定的神经元是否重要，但却忽略了这个神经元是如何与大脑其他部分进行交流的。
在剪枝的同时训练模型： 它们试图在缩小模型的同时重新教它如何思考。这就像是在重建发动机的同时学习驾驶汽车。这既耗时漫长，又需要消耗大量的燃料（计算能力）。

解决方案：TraceNAS（“梯度轨迹”侦探）

作者提出了一种名为 TraceNAS 的新方法。你可以把它想象成一个高科技 X 光扫描仪，它可以在不实际接触或重新教导模型的情况下，观察模型的“灵魂”。

以下是它的工作原理，使用简单的类比说明：

1. 原型思维的“回声”

想象原始的巨型模型是一位烹饪出完美佳肴的大厨。当你品尝时，味道会在你的舌尖留下特定的“痕迹”。

旧方法： 为了看一个新的小食谱是否可行，你必须烹饪整道菜，品尝它，如果不好吃，就得从头开始。
TraceNAS 的方法： TraceNAS 不去烹饪整道菜，而是通过观察食材和烹饪步骤（梯度）来预测味道是否会正确。它检查缩小后的版本的“味道痕迹”是否与原始大厨的“味道痕迹”相匹配。

2. “影子”测试（梯度轨迹相关性）

论文使用了一个数学概念——梯度轨迹相关性（Gradient Trace Correlation）。

想象原始模型是一艘在平静海洋上航行的巨轮。它在身后留下了一道特定的航迹（波浪轨迹）。
当你尝试建造一艘较小的船（被剪枝后的模型）时，TraceNAS 会询问：“这艘小船留下的航迹看起来是否与大船的航迹完全一致？”
如果航迹匹配，这意味着小船正沿着相同的路径行驶，并且一旦经过一点点练习，就能达到相同的目的地（表现出色）。如果航迹混乱，那么这艘船将会翻船。

3. “低秩”捷径

为巨轮计算这些航迹通常需要超级计算机。TraceNAS 非常聪明：它意识到巨轮的运动主要发生在几个主要方向上。它利用**低秩（Low-Rank）**技巧，只观察航迹中最重要的几个方向。

类比： 与其测量海洋中的每一个涟漪，它只测量三个最大的波浪。这使得它们仅用一块显卡（GPU）在 8.5 小时内就能完成测试，而旧方法则需要在整个计算集群上花费数天甚至数周的时间。

结果：快速、廉价且智能

论文在 Llama 和 Qwen 等知名模型上测试了这种方法。

速度： 他们发现寻找最佳“剪切点”仅需在单台电脑上花费 8.5 小时。而旧方法比它们慢 10 倍，且消耗了 10 倍的能量。
准确度： 最终生成的轻量化模型表现得与那些为了寻找最佳剪切点而进行了数周训练的模型一样出色。
非均匀性： 与那个“剪掉一切都一样”的“笨拙园丁”不同，TraceNAS 找到了“定制化的契合度”。它保留了大脑中沉重且复杂的关键部分，只修剪了脂肪，从而创造出一个高效且依然非常聪明的模型。

核心结论

TraceNAS 是一个工具，它能让你在不需要重新教导模型或花费巨额电费的情况下，将一个巨大的 AI 大脑缩小成口袋大小的版本。它通过检查小脑的“影子”是否与大脑的“影子”相匹配来实现这一点。如果影子对齐了，说明这个小脑已经准备好了。

这使得在普通设备上运行强大的 AI 成为可能，而无需依赖庞大的数据中心，同时也节省了大量的时间和能源。

技术摘要：TraceNAS

问题定义

大语言模型（LLMs）的部署受到庞大的参数量和计算需求的限制。虽然结构化剪枝通过移除架构单元（如注意力头、Transformer 块）为高效部署提供了一条路径，但现有方法面临显著局限性：

局部与全局依赖关系： 大多数剪枝方法使用局部启发式算法（如激活加权幅度或基于 Hessian 的重构）在孤立的状态下评估层、头或通道的重要性。这些方法忽略了复杂的全局结构依赖以及剪枝引起的误差在模型中的传播。
训练感知搜索的计算成本： 试图捕捉全局依赖关系的方法通常依赖于训练感知的神经架构搜索（NAS）。这些方法需要迭代的梯度更新和持续的未见标记流以防止过拟合，这使得搜索过程的计算成本与剪枝后的恢复训练本身一样高。
现有零样本代理（Zero-Shot Proxies）的不足： 当前的无训练代理旨在通过关注可训练性（如梯度稳定性、损失平面平滑度）来识别可训练的随机初始化架构。它们无法评估“功能继承”（Functional Inheritance）——即剪枝后的模型通过保留预训练损失平面来恢复性能的能力。

方法论

TraceNAS 被提议为一个旨在无需搜索时训练即可联合优化 LLM 深度和宽度的无训练 NAS 框架。

1. 搜索空间与编码

该框架将预训练的 LLM 视为一个超网络。搜索空间 $G$ 被定义为以下各项的联合分布：

深度 ( $d$ )： $L$ 个 Transformer 块上的离散掩码，允许在保持残差流连续性的同时进行非均匀深度缩减。
宽度 ( $\kappa$ )： 注意力子块和 MLP 子块的参数保留比例，实现模型内异构的稀疏性。

2. 原地架构实现

为了避免实例化数千个候选模型的内存开销，TraceNAS 采用了原地掩码策略（In-Place Masking Strategy）：

宽度掩码： 使用激活加权乘积（类似于 Wanda）为权重通道生成二进制掩码。这种 $O(d^2)$ 启发式方法识别功能关键通道，同时保留离群特征。
深度掩码： 如果被停用，则跳过整个块，通过残差连接路由激活。
执行： 框架临时修改基础模型的权重矩阵以在极小的校准集上计算梯度，然后恢复原始权重，从而为后续迭代保留基础模型。

3. 零样本代理：梯度轨迹相关性

其核心创新是一个标度不变的代理 $\Phi$ ，用于衡量剪枝后的候选模型 ( $M_{sub}$ ) 与预训练基础模型 ( $M_{base}$ ) 之间的功能对齐（Functional Alignment）。

梯度轨迹： 在校准集上计算梯度。为确保可扩展性，这些梯度被投影到低秩（LoRA）子空间中，利用 LLM 内在的低维特性，将内存开销降低了数个数量级。
对齐指标： 代理计算每一层标准化梯度轨迹的皮尔逊相关系数（Pearson Correlation Coefficient）。这确保了标度不变性，侧重于方向对齐而非幅度偏移。
稀疏权重聚合： 最终得分 $\Phi$ 是层级相关性的加权总和，其中权重对应于参数保留比例。这使全局得分锚定在高容量区域，防止来自过度剪枝的子块的噪声扭曲排名。

4. 进化搜索

TraceNAS 采用进化策略在离散搜索空间中进行导航：

初始化： 基于全局重要性先验生成种群。
进化： 种群经历交叉和变异。深度进化使用位翻转（Bit-flip）变异和连续范围交叉，而宽度进化使用插值和高斯抖动。
选择： 候选模型仅根据代理 $\Phi$ 进行排名，选择那些在剪枝后持续预训练（CPT）期间具有最高性能恢复潜力的模型。

核心贡献

可扩展的无训练 NAS： TraceNAS 提供了一个统一的框架，用于联合深度和宽度剪枝，能够在无需搜索时训练的情况下发现非均匀架构。
功能继承代理： 引入了一种新型的、标度不变的基于梯度的代理，用于评估剪枝模型与预训练模型之间的梯度轨迹对齐情况。它实现了极高的下游准确率相关性（Spearment $\rho = 0.94$ , Kendall $\tau = 0.82$ ）。
高效性： 该框架通过利用低秩梯度子空间和原地掩码来减少搜索开销，实现在单张 GPU 上进行高保真度的模型发现。

实验结果

TraceNAS 在 Llama-2-7B、Llama-3.1-8B 和 Qwen-2.5-14B 上进行了评估，分别将其剪枝至 2.7B、4.6B 和 8.4B 参数。

搜索效率： TraceNAS 在 8.5 个 GPU 小时内（在 A100 上）识别出最优架构，与 DarwinLM 等训练感知基准相比，搜索时间减少了 10 倍。它仅需 98M 搜索标记，比 DarwinLM 少了 10 倍。
性能：
- Llama-2-7B： TraceNAS 在八个推理基准测试中实现了最高的平均准确率 (62.81%)，超过了 DarwinLM (62.57%) 和 ShearedLLaMA (62.63%)。值得注意的是，当 ShearedLLaMA 受限于 TraceNAS 所使用的相同标记预算时，其性能显著下降（降至 60.86%）。
- Llama-3.1-8B & Qwen-2.5-14B： 在复杂推理任务（如 ARC-Challenge, BoolQ）上，TraceNAS 优于均匀剪枝基准和训练感知方法（DarwinLM, E3-Pruner），且所需的恢复标记显著更少（20B vs. Minitron 的 94B）。
代理验证： 该代理在不同超参数（样本量、上下文长度和 LoRA 秩）下表现出高度稳定性，证实了梯度轨迹对齐即使在受限设置下也能有效捕捉模型继承性。

意义与主张

论文声称 TraceNAS 解决了结构化剪枝中计算成本的关键瓶颈。通过将重点从“可训练性”（从零开始学习）转向“继承性”（保留预训练损失平面），TraceNAS 能够以极低的资源消耗发现高性能的非均匀剪枝模型。

作者强调，其方法：

消除了昂贵的搜索时训练成本，使大规模模型压缩变得触手可及。
捕捉了局部启发式算法会忽略的全局结构依赖，确保了复杂推理能力的保留。
降低了模型压缩的碳足迹，通过大幅降低架构搜索所需的 GPU 小时数来实现。

这项工作将 TraceNAS 定位为在资源受限的环境中部署高效 LLM 的可扩展基础，同时保持密集预训练模型的函数完整性。