Source-linked AI summary
Attention Pruning: Automated Fairness Repair of Language Models via Surrogate Simulated Annealing
Vishnu Asutosh Dasu, Md Rafi ur Rashid, Vipul Gupta, Saeid Tizpaz-Niari, Gang Tan
TL;DR
Fairness-aware attention-head pruning must balance bias reduction against utility, while searching combinations of heads is computationally challenging. The paper uses surrogate DNNs with simulated annealing to search pruning configurations. It reports up to 40% reduction in gender bias and better fairness and utility than state-of-the-art methods in 4/6 cases.
Problem
Finding attention-head subsets that reduce bias while preserving utility is a combinatorial problem, and head interactions can make individual-head assessments insufficient.
Source: §1 · p. 1, §3.2 · pp. 3–4Method
The method formulates head pruning as discrete optimization and uses surrogate DNNs to predict bias and perplexity during simulated annealing.
Source: §4.1 · p. 4, §4.2 · p. 5Results
AP reduced gender bias by up to 40% and found better fairness and utility than state-of-the-art methods in 4/6 cases.
Source: §1 · p. 2, §8 · p. 11Takeaways & Limitations
The authors conclude that surrogate simulated annealing is an effective, scalable strategy for improving LLM fairness while preserving utility.
Source: §8 · p. 11Exploring the top k search states with genetic algorithms, rather than selecting the best state, remains future work.
Source: §8 · p. 11
Abstract
from arXiv · showhide
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.
1 Introduction
The introduction frames fairness-aware pruning as a combinatorial search problem: attention-head combinations can interact nonlinearly, while fairness and utility must be balanced. Attention Pruning uses surrogate DNNs with simulated annealing and reports reduced bias and improved fairness-utility outcomes across several LLMs and bias categories.
1 Introduction
- The introduction situates the work amid prior fairness-repair research and notes that repairing unfairness in frontier AI models has been understudied in software engineering.It also describes the fairness-utility trade-off as an ongoing challenge in language-model bias mitigation.Source: §1 · p. 1
- Pruning is a combinatorial search over attention-head subsets, and head combinations can have nonlinear effects on fairness and utility.Prior fairness-aware pruning analyzes individual heads, which does not account for these interactions.Source: §1 · p. 1
- Surrogate DNNs approximate the fairness and utility of pruned LLMs, allowing simulated annealing to search without repeatedly evaluating the billion-parameter models.The introduction identifies high inference time as the barrier to directly applying simulated annealing.Source: §1 · p. 2
- Unfairness fell by up to 40%, and AP found better fairness and utility than state-of-the-art methods in 4/6 cases.The authors also report fairness improvements of up to 65% for race, nationality, sexual-orientation, and age bias.Source: §1 · p. 2
- The authors report that surrogate DNNs predict pruning effects with MSE ≤0.01 and that a hyperparameter can explore the fairness-utility Pareto frontier.They describe this hyperparameter as a way to tune the trade-off between fairness and utility.Source: §1 · p. 2
- Reducing gender bias with AP also reduced other social biases, with fairness improvements of up to 65% across the reported evaluations.The evaluated categories included race, nationality, sexual orientation, and age.Source: §1 · p. 2
2 Background
The background defines causal language modeling and describes bias types and evaluation. It presents perplexity as a utility measure and HolisticBias toxicity differences across subgroups as a way to quantify fairness.
2.1 Causal Language Modeling
- In causal language modeling, a model learns a next-token distribution, minimizes negative log-likelihood, and generates text autoregressively until a length limit or STOP token.The background describes the objective over a sequence of tokens.Source: §2.1 · p. 2
2.2 Bias in Causal Language Modeling
- Bias in causal language models includes intrinsic bias from data, algorithms, or architecture and extrinsic bias measured in task-specific applications.The background gives gender stereotypes as an intrinsic-bias example and describes application-specific benchmarks for extrinsic bias.Source: §2.2 · p. 2
2.3 Notions of Model Utility via Perplexity
- The paper measures language-model utility with perplexity, computed from causal language-modeling loss on a fixed token sequence and evaluated on WikiText-2.Lower perplexity indicates a model better estimates the real-world probability distribution.Source: §2.3 · p. 2
2.4 Notions of Model Fairness
- HolisticBias evaluates toxicity on prompts across social groups; the paper quantifies bias through differences in subgroup toxicity, so low average toxicity alone does not imply low bias.The dataset contains 566,000 prompts in 13 social-bias categories, and toxicity is scored by a pretrained BERT classifier.Source: §2.4 · pp. 2–3
3 Problem Statement
The section frames fairness-aware attention-head pruning as selecting heads to remove so bias is minimized with the least degradation in model utility. It describes the Transformer attention structure and contrasts FASP’s individual-head ranking with the nonlinear interactions among heads.
3.1 Multi-Head Attention in Transformers
- Transformer blocks pair a Multi-Head Attention layer with a fully connected layer; attention heads process queries, keys, and values derived from earlier outputs.The first block derives these inputs from text embeddings.Source: §3.1 · p. 3
3.2 Attention Head Pruning for Fairness
- Structured attention-head pruning removes whole heads, which are omitted from concatenation before the output projection.The section contrasts this with unstructured pruning at the neuron level.Source: §3.2 · p. 3
- Fairness-aware pruning selects a subset of attention heads to remove, minimizing bias while trading off utility loss measured by perplexity.The objective weights bias and perplexity loss using ε.Source: §3.2 · p. 3
- FASP ranks heads by perplexity impact before ranking the remaining heads by bias impact, but it does not model nonlinear interactions among heads.The authors note that individually beneficial removals may jointly fail to reduce bias and may worsen perplexity.Source: §3.2 · p. 4
4 Approach
Attention Pruning uses surrogate-assisted simulated annealing to select attention heads for pruning while balancing bias and perplexity. It models nonlinear head interactions and evaluates the selected configuration on the LLM.
4.1 Discrete Optimization with Simulated Annealing
- SA searches a 2^N head-subset space, balancing bias and perplexity in its cost while temperature controls exploration versus exploitation.Source: §4.1 · pp. 4–5
- Bounds n_l and n_u restrict the number of pruned heads, avoiding too few or too many pruned heads and reducing the search space.Source: §4.1 · p. 4
4.2 Approximating Bias and Perplexity of Pruned Large Language Models
- Surrogate DNNs are trained by regression on randomly sampled pruned-model evaluations using HolisticBias prompts and WikiText-2 text.The training objective minimizes mean squared error.Source: §4.2 · p. 5
4.3 Efficient Surrogate Simulated Annealing for Pruning Attention Heads
- Surrogate-based SA minimizes predicted bias and perplexity over pruning states, then prunes the selected heads and computes the actual metrics.Source: §4.3 · p. 5
- ≈0.34 milliseconds on a CPU for surrogate inference on Llama-2-7B with 1024 heads enables ≈2940 states per second, a ≈2,260,000× speed-up over vanilla SA.The reported vanilla SA comparison requires expensive GPU resources.Source: §4.3 · p. 5
5 Experiments
Attention Pruning (AP) uses surrogate simulated annealing to search for attention heads to prune, balancing gender-bias reduction with model perplexity. Across six LLMs, it generally outperforms fairness-aware pruning baselines and also reduces other evaluated social biases.
Experimental setup
- Across six LLMs, experiments focus on HolisticBias gender bias and compare AP with FASP, magnitude-based pruning, and random pruning.The authors use an exhaustive search over pruning ratios for FASP and AP, and limit the simulated-annealing search to three hours.Source: §5 · pp. 5, 7, 8
Surrogate DNN training
- Small two-hidden-layer surrogate DNNs predict bias and perplexity; the authors collect at least 25,000 samples per model and tune hidden-layer sizes by model scale.They use 64 and 32 neurons for smaller LLMs and 256 and 128 for larger LLMs.Source: §5 · p. 6
Dataset preparation
- Gender-bias prompts are subsampled to preserve subgroup proportions, while perplexity targets are clamped before scaling to address high variance without discarding high-perplexity states.The authors set the standard-deviation threshold for clamping at σ ≤ 10.Source: §5 · p. 7
RQ1: Surrogate effectiveness
- The surrogate DNNs have low validation MSE; using only the bias or perplexity predictor favors that objective, while their combined cost yields states with low bias and perplexity.The reported pattern is high perplexity with bias-only optimization and high bias with perplexity-only optimization.Source: Table 4 · p. 7, §5 · pp. 7–8, Table 5 · p. 8
RQ2: Comparison with pruning strategies
- Overall, AP reduces bias more while ensuring better perplexity than FASP in most cases; for GPT-J-6B, the greatest absolute and relative bias improvements are 0.18 and 40.8%.The other strategies, which do not account for bias, are worse than both AP and FASP.Source: §5 · p. 8
RQ3: Design considerations
- Larger validation subsets can improve surrogate estimates; increasing ε prioritizes bias reduction at perplexity's expense, and longer SA runs may find better solutions.The authors report that bias and perplexity can alternate as the search runs longer.Source: §5 · pp. 8–9, Figure 2 · p. 8, Figure 3 · p. 9
RQ4: Effects on other social biases
- After gender-bias repair, all four additionally evaluated social biases decrease across the LLMs, with reductions of up to 65%.The highest reported absolute improvement is 0.29 for sexual-orientation bias on GPT-Neo-125M; the highest relative improvement is 65.21% for age bias on GPT-J-6B.Source: Table 7 · p. 9, §5 · pp. 9–10
RQ4: Generalization to racial bias
- AP also reduces racial bias when optimizing for it directly; the authors report better results from direct racial-bias reduction than from reducing it via gender-bias repair.This direct optimization is distinct from evaluating how gender-bias repair affects other biases.Source: §5 · pp. 9–10, Table 8 · p. 10
6 Related Works
The related-work discussion places AP among fairness-aware pruning and broader bias-mitigation approaches. It contrasts direct search over LLM configurations with methods that alter data, training, or inference behavior.
Fairness-aware pruning and repair
- Prior work on fairness-aware pruning includes FASP, while NeuFair uses simulated annealing to drop neurons but does not scale to LLMs because it evaluates the model directly at each search step.The authors note that NeuFair's experiments were limited to small DNNs with at most 384 neurons.Source: §6 · p. 10
Bias-mitigation approaches
- Bias mitigation spans pre-processing, in-processing, and post-processing, including data curation, augmentation, model regularization, and inference-time interventions.Examples include counterfactual data generation, adapter layers, token blocking, and sparse debiasing subnetworks.Source: §6 · p. 10
Motivation for post-processing
- The authors frame post-processing as useful when dataset collection and training are expensive or may not be repeatable for large pre-trained LLMs.They describe post-processing as sometimes the only viable option given the resources required to train these models.Source: §6 · p. 10
7 Discussion
The discussion reports AP’s bias and perplexity advantages over FASP and illustrates subgroup toxicity differences. It also notes limits from surrogate/search randomness and the scope of the evaluation.
Practical Significance of Bias Reduction
- AP usually improves bias over FASP by ≈0.02 and perplexity by up to ≈6 points or ≈11%.The comparison covers all 33,789 test prompts and 8 Gender-group subgroups in HolisticBias.Source: §7 · p. 10
- For GPT-J-6B, Cliff’s δ = −0.968 (95% CI ≈[−0.994, −0.847]); for Llama-2-7B, δ = −0.9392 (95% CI ≈[−0.984, −0.785]), both indicating large effects.The reported one-sided Mann–Whitney p-values are 2.319 × 10^-9 for GPT-J-6B and 6.539 × 10^-9 for Llama-2-7B.Source: §7 · pp. 10–11
Qualitative Examples
- 40-token examples show AP and FASP reduce toxicity differences between sex and cisgender prompts, but neither guarantees zero toxicity.The authors describe the baseline as producing more toxic content for the sex subgroup than for the cisgender subgroup.Source: §7 · p. 10, Table 9 · p. 11
Transferability of Surrogate DNNs
- The surrogate-DNN idea extends to Race bias, but the tested DNN architecture is not intended to transfer across datasets and fairness metrics.Architecture and layer choices depend on the bias and perplexity datasets used for training.Source: §7 · p. 11
Limitations
- Surrogate prediction and simulated annealing add randomness and can yield suboptimal results; collecting 27,000 Llama-2-7B samples and training the surrogates took ≈1,900 GPU-hours.The authors attribute one GPT-Neo-125M outcome to possible DNN prediction error and note that simulated annealing may require infeasible iterations to reach an optimum.Source: §7 · p. 11
Threats to Validity
- Evaluation covered six LLMs but only HolisticBias and WikiText-2, so generalization to all models and bias metrics remains unknown.The authors cite limited computational resources as the reason for this dataset scope.Source: §7 · p. 11
8 Conclusion
The conclusion presents attention-head pruning with surrogate simulated annealing as a scalable way to improve LLM fairness while preserving utility. AP achieves up to 40% gender-bias reduction and outperforms state-of-the-art post-processing bias mitigation; the authors identify other responsible-AI applications as future work.
8 Conclusion
- Surrogate simulated annealing is presented as an effective, scalable strategy for the computationally hard task of pruning attention heads while preserving utility.AP selects an ideal subset of heads to prune as a post-processing fairness repair.Source: §8 · p. 11
- Up to 40% reduction in gender bias; AP also outperforms state-of-the-art approaches with attention-head configurations achieving both low bias and low perplexity.The conclusion also reports that AP outperforms the state-of-the-art post-processing bias mitigator.Source: §8 · p. 11
- The authors propose exploring top-k search states with genetic algorithms and extending head pruning to privacy and robustness as future work.The proposed genetic-algorithm direction would explore the top-k states rather than select only the best state.Source: §8 · p. 11