Attention Pruning: Automated Fairness Repair of Language Models via Surrogate Simulated Annealing

📄 arXiv: 2503.15815v3 📥 PDF

作者: Vishnu Asutosh Dasu, Md Rafi ur Rashid, Vipul Gupta, Saeid Tizpaz-Niari, Gang Tan

分类: cs.AI

发布日期: 2025-03-20 (更新: 2025-11-25)

💡 一句话要点

提出Attention Pruning，通过代理模拟退火自动修复语言模型中的偏见。

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

关键词: 语言模型 公平性 偏差缓解 注意力剪枝 模拟退火 代理模型 后处理 程序修复

📋 核心要点

  1. 大型语言模型存在偏见，但修改训练数据或算法成本高昂，后处理剪枝是一种可行的替代方案。
  2. 提出Attention Pruning方法，使用代理模型加速搜索过程，通过模拟退火算法寻找最优的注意力头剪枝方案。
  3. 实验表明，Attention Pruning能够有效降低性别偏见，最高可达40%，并且优于其他现有的偏差缓解方法。

📝 摘要（中文）

本文探讨了剪枝注意力头作为大型语言模型（LLM）后处理偏差缓解方法。LLM等现代AI系统正扩展到敏感的社会环境中，公平性问题变得尤为重要。由于LLM通过在大量人类生成内容的数据集上训练来发展决策模式，因此它们自然会编码并延续社会偏见。修改训练数据集和算法成本高昂且需要大量资源；而后处理技术（例如选择性地停用预训练LLM中的神经元和注意力头）可以提供可行且有效的方法来提高公平性。然而，识别要修剪的最佳参数子集在LLM的巨大参数空间中提出了一个组合挑战，需要能够有效平衡模型公平性和效用目标的解决方案。为了应对计算挑战，我们探索了一种通过随机模拟退火的基于搜索的程序修复方法。鉴于数十亿参数LLM中令人望而却步的评估成本，我们开发了代理深度神经网络，可以有效地建模注意力头状态（激活/非激活）与其相应的公平性/效用指标之间的关系。这使我们能够对代理模型执行优化，并有效地识别注意力头的最佳子集以进行选择性修剪，而不是直接搜索LLM参数空间。本文介绍了一种公平性感知的代理模拟退火方法Attention Pruning，用于修剪LLM中不成比例地导致偏差的注意力头，同时最大限度地减少对整体模型效用的影响。我们的实验表明，Attention Pruning在性别偏见方面实现了高达40％的降低，并且优于最先进的偏差缓解策略。

🔬 方法详解

问题定义：大型语言模型（LLM）在训练过程中会学习并放大社会偏见，导致在下游任务中产生不公平的结果。现有缓解偏见的方法，如修改训练数据或调整训练算法，通常需要大量的计算资源和时间，并且可能影响模型的整体性能。因此，如何在不重新训练模型的情况下，高效地减少LLM中的偏见是一个关键问题。

核心思路：本文的核心思路是通过剪枝LLM中的注意力头来减少偏见。具体来说，作者认为某些注意力头可能对模型产生偏见起着更大的作用，因此选择性地移除这些注意力头可以在降低偏见的同时，尽量保持模型的性能。为了解决在巨大的参数空间中搜索最优剪枝方案的计算挑战，作者引入了代理模型来加速搜索过程。

技术框架：Attention Pruning方法包含以下几个主要阶段： 1. 代理模型训练：使用深度神经网络来模拟注意力头状态（激活/非激活）与公平性/效用指标之间的关系。代理模型的目标是快速预测给定注意力头配置下的模型性能和公平性。 2. 模拟退火优化：使用模拟退火算法在代理模型上搜索最优的注意力头剪枝方案。模拟退火算法是一种随机搜索算法，可以在搜索空间中探索不同的解决方案，并逐步收敛到最优解。 3. 注意力头剪枝：根据模拟退火算法找到的最优方案，对LLM进行注意力头剪枝。

关键创新：Attention Pruning的关键创新在于使用代理模型来加速注意力头剪枝的搜索过程。与直接在LLM上进行搜索相比，代理模型可以显著降低计算成本，从而使得在大规模LLM上进行偏差缓解成为可能。此外，使用模拟退火算法可以有效地探索巨大的参数空间，找到在公平性和效用之间取得良好平衡的剪枝方案。

关键设计： * 代理模型结构：作者使用深度神经网络作为代理模型，输入是注意力头状态（激活/非激活），输出是公平性指标（如性别偏见）和效用指标（如模型准确率）。 * 损失函数：代理模型的损失函数包括公平性损失和效用损失，通过调整两个损失的权重来平衡公平性和效用。 * 模拟退火参数：模拟退火算法的关键参数包括初始温度、降温速率和迭代次数。这些参数需要根据具体问题进行调整，以获得最佳的搜索效果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Attention Pruning在降低性别偏见方面取得了显著效果，最高可降低40%。与现有的偏差缓解方法相比，Attention Pruning在降低偏见的同时，对模型性能的影响更小。此外，代理模型的使用显著降低了计算成本，使得该方法能够应用于大规模LLM。

🎯 应用场景

Attention Pruning可应用于各种需要公平性保障的自然语言处理任务，例如简历筛选、信用评估、新闻推荐等。通过降低LLM中的偏见，可以避免模型在这些应用中产生歧视性结果，从而提高社会公平性。该方法还可用于提升模型的可解释性，帮助人们理解模型决策过程中的潜在偏见来源。

📄 摘要（原文）

This paper explores pruning attention heads as a post-processing bias mitigation method for large language models (LLMs). Modern AI systems such as LLMs are expanding into sensitive social contexts where fairness concerns become especially crucial. Since LLMs develop decision-making patterns by training on massive datasets of human-generated content, they naturally encode and perpetuate societal biases. While modifying training datasets and algorithms is expensive and requires significant resources; post-processing techniques-such as selectively deactivating neurons and attention heads in pre-trained LLMs-can provide feasible and effective approaches to improve fairness. However, identifying the optimal subset of parameters to prune presents a combinatorial challenge within LLMs' immense parameter space, requiring solutions that efficiently balance competing objectives across the frontiers of model fairness and utility. To address the computational challenges, we explore a search-based program repair approach via randomized simulated annealing. Given the prohibitive evaluation costs in billion-parameter LLMs, we develop surrogate deep neural networks that efficiently model the relationship between attention head states (active/inactive) and their corresponding fairness/utility metrics. This allows us to perform optimization over the surrogate models and efficiently identify optimal subsets of attention heads for selective pruning rather than directly searching through the LLM parameter space. This paper introduces Attention Pruning, a fairness-aware surrogate simulated annealing approach to prune attention heads in LLMs that disproportionately contribute to bias while minimally impacting overall model utility. Our experiments show that Attention Pruning achieves up to $40\%$ reduction in gender bias and outperforms the state-of-the-art bias mitigation strategies.