ALPS: Attention Localization and Pruning Strategy for Efficient Alignment of Large Language Models

本文提出 ALPS 算法，通过基于权重分布的参数对齐分布分数（sPAD）定位任务敏感注意力头并剪枝，仅更新 10% 的注意力参数即在通用、数学和代码任务上实现性能提升，同时展现头部可转移性和知识遗忘缓解效果。

Large Language Model, Fine-tuning, Efficiency, Parameter-Efficient Fine-Tuning, Reasoning

Hao Chen, Haoze Li, Zhiqing Xiao, Lirong Gao, Qi Zhang, Xiaomeng Hu, Ningtao Wang, Xing Fu, Junbo Zhao

Zhejiang University, Ant Group

Generated by grok-3

Background Problem

大型语言模型（LLMs）在各种任务中表现出色，但将其对齐到下游任务通常需要构建任务特定的指令数据集和进行广泛的训练调整，导致资源消耗巨大。近年来，研究者尝试通过最小化数据需求或基于数据激活来识别关键注意力头以提高对齐效率，但这些方法引入了数据依赖性，限制了泛化性和可重用性。本文提出了一种新方法，旨在通过利用模型权重参数来定位任务敏感注意力头，解决对齐过程中的效率问题，同时避免数据依赖性，并探索注意力头在任务对齐中的作用。

Method

ALPS（Attention Localization and Pruning Strategy）是一种基于启发式搜索的算法，旨在提高大型语言模型对齐下游任务的效率，其核心思想和步骤如下：

• 核心思想：通过分析模型权重参数的分布变化，识别对下游任务最敏感的注意力头，并通过限制训练更新到这些头部来减少计算成本，同时避免数据依赖性。
• 具体实现：
  1. 注意力头定位：给定预训练模型 $M^B$ 和任务微调模型 $M^T$，提取每个注意力头的投影矩阵 $W_o^h$，并通过 softmax 函数将其转换为概率分布 $P^h$。然后，使用 Wasserstein-1 距离计算基础模型和任务模型之间分布的偏移，定义为参数对齐分布分数（sPAD），即 $s_h^{PAD} = W_1(P_B^h, P_T^h)$，用于量化每个头部的任务敏感性。
  2. 注意力头剪枝：根据 sPAD 分数选择 Top-K 个任务敏感注意力头（默认比例为 10%），在微调过程中冻结其余头部的梯度更新，即对非敏感头部设置 $abla_{ heta^h} \mathcal{L} = 0$，从而减少优化冗余。
• 关键点：ALPS 不依赖任务特定数据激活参数，而是直接从权重矩阵中提取任务相关信息，试图提高泛化性和可重用性。然而，sPAD 度量的选择（基于 Wasserstein-1 距离）缺乏充分的理论依据，是否能准确捕捉任务相关性尚待验证。此外，冻结大部分注意力头可能导致模型对复杂任务的适应能力不足，特别是在任务需求多样化时。

Experiment

实验在 Llama-3 系列模型（1B、3B、8B 参数规模）上进行，覆盖通用、数学和代码生成三个下游任务，使用的数据集包括 UltraChat、MathInstruct 和 MagiCoder，并额外测试了 Alpaca、Camel-math 和 CodeAlpaca 以评估可转移性。评估指标通过 lm-eval 和 EvalPlus 工具实现，涵盖指令跟随、推理、数学能力和代码生成等多个方面。

• 实验设置：与多种基线方法对比，包括全参数微调（w/ full）、冻结所有注意力头（w/o attn）、随机选择头部（Random）、层一致性选择（LC）和 LoRA 方法。ALPS 默认选择 10% 的注意力头进行更新。
• 结果分析：ALPS 在所有模型规模和任务上均优于基线方法，例如在 Llama-3.2-1B 上平均性能提升 2.75%，在 Llama-3.1-8B 上提升 1.25%。消融研究表明，10% 和 30% 的头部选择比例在性能和效率之间取得最佳平衡。此外，ALPS 识别的任务敏感头部在同一任务域内不同数据集上表现出可转移性，并在通用基准（如 MMLU 和 ARC-C）上缓解了知识遗忘问题。
• 评价与质疑：虽然结果显示 ALPS 有效提升了效率和性能，但实验设计存在局限性：数据集选择较为狭窄，未充分覆盖任务分布的多样性；可转移性测试仅限于同一任务域，跨任务域的泛化能力未被验证；此外，sPAD 度量与其他度量（如 KL 散度）的对比虽显示优势，但缺乏对为何 Wasserstein-1 距离更适合的深入分析。实验结果的稳健性和方法在更大规模或更复杂任务上的适用性仍需进一步验证。

Further Thoughts

ALPS 的核心创新在于利用权重分布变化来定位任务敏感注意力头，这一思路为参数高效微调提供了新视角。然而，方法中对 sPAD 度量的依赖引发了一些思考：是否可以通过结合其他分布距离度量（如 Jensen-Shannon 散度）或引入任务无关的先验知识来进一步提高识别精度？此外，论文未探讨注意力头在不同任务间的潜在冲突，例如数学任务和代码任务可能需要共享某些注意力头，但其功能需求可能存在差异，如何在这种情况下优化头部选择策略是一个值得深入研究的方向。

另一个有趣的点是 ALPS 在缓解知识遗忘方面的潜力。论文虽提供了初步证据，但未深入探讨其机制。结合近期关于模型稀疏性和过拟合的研究，或许可以通过分析注意力头的稀疏更新如何影响模型的记忆容量来进一步解释这一现象。此外，ALPS 的方法可能与其他参数高效微调技术（如 LoRA 或 Adapter）结合，通过在任务敏感头部上应用低秩更新，进一步减少计算成本并提升性能。这种结合可能在资源受限的场景（如边缘设备上的模型部署）中具有实际应用价值。

最后，ALPS 的可转移性测试局限于同一任务域内，未来研究可以探索跨任务域甚至跨模态任务的注意力头共享机制，例如是否可以将数学任务中识别的注意力头应用于科学计算任务？这可能为构建更通用、更高效的基础模型对齐框架铺平道路。