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

This paper explores pruning attention heads as a post-processing bias mitigation method for large language models (LLMs). As LLMs expand into sensitive social contexts, fairness concerns become crucial since these models encode and perpetuate societal biases from their training data. While modifying training datasets is prohibitive, post-processing techniques like pruning attention heads can provide feasible approaches to improve fairness. However, identifying optimal parameters to prune presents a combinatorial challenge requiring efficient solutions that balance model fairness and utility. We introduce Attention Pruning, a fairness-aware surrogate simulated annealing approach that uses surrogate deep neural networks to efficiently model the relationship between attention head states and fairness/utility metrics, enabling optimization without directly searching the massive LLM parameter space. Our experimental evaluation shows that Attention Pruning achieves up to \(40\%\) reduction in gender bias, outperforming the state-of-the-art.