ALPS: Attention Localization and Pruning Strategy for Efficient Adaptation of Large Language Models
Authors: Hao Chen, Haoze Li, Zhiqing Xiao, Lirong Gao, Qi Zhang, Xiaomeng Hu, NINGTAO WANG, Xing Fu, Junbo Zhao
Accepted at: acl25

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

Hao Chen*, Haoze Li*, Zhiqing Xiao*, Lirong Gao*, Qi Zhang*
Xiaomeng Hu*, Ningtao Wang*, Xing Fu*, Junbo Zhao*
*Zhejiang University   *Ant Group
{h.c.chen, j.zhao}@zju.edu.cn

Abstract

Aligning general-purpose large language models (LLMs) to downstream tasks often incurs significant training adjustment costs. Prior research has explored various avenues to enhance alignment efficiency, primarily through minimal-data training or data-driven activations to identify key attention heads. However, these approaches inherently introduce data dependency, which hinders generalization and reusability. To address this issue and enhance model alignment efficiency, we propose the Attention Localization and Pruning Strategy (ALPS), an efficient algorithm that localizes the most task-sensitive attention heads and prunes by restricting attention training updates to these heads, thereby reducing alignment costs. Experimental results demonstrate that our method activates only \(10\%\) of attention parameters during fine-tuning while achieving a \(2\%\) performance improvement over baselines on three tasks. Moreover, the identified task-specific heads are transferable across datasets and mitigate knowledge forgetting. Our work and findings provide a novel perspective on efficient LLM alignment. The code is available at https://github.com/VoiceBeer/ALPS.

1 Introduction

Pre-trained Large Language Models (LLMs) have demonstrated impressive performance across various tasks (Achiam et al., 2023; OpenAI, 2025; Guo et al., 2025). A way to harness the potential of pre-trained foundation models on downstream tasks is task alignment, which mainly incorporates tailored task knowledge. Recently, task alignment has garnered considerable attention in both industry and academia with a diverse range of specialized models having been developed, spanning domains such as mathematics (Shao et al., 2024; Yang et al., 2024), code generation (Guo et al., 2024; Roziere et al., 2023), medical diagnostics (Singhal et al., 2023, 2025) and protein structure prediction (Schmirler et al., 2024). However, despite the effectiveness of task alignment, it still poses a significant resource-intensive challenge. The alignment process demands extensive efforts to construct task-specific instruction datasets and training adjustments, and further, its computational demands exacerbate the costs (Zhao et al., 2023). These challenges highlight the pressing need for more efficient alignment strategies that not only minimize resource consumption costs but also maintain or even enhance model performance (Wan et al., 2023).

To address the efficiency challenges of model alignment, recent LLM studies are interested in attention mechanisms and attempt to localize task-specific attention heads (Clark, 2019; Michel et al., 2019; Zhou et al., 2024b; Wu et al., 2024; Tang et al., 2024). However, these methods mainly rely on task-specific data to activate model parameters for attention head localization, inadvertently introducing data dependency to identified heads, which not only hinders generalization but also complicates reusability. Moreover, these approaches over-

![Process of localizing task-sensitive attention heads. The left side represents the attention heads of LLM, and the center illustrates the distribution of head parameters, where task parameters \( \theta_T \) capture task information and lead to task-sensitive heads \( H_T \).](page_682_370_627_370.png)

Figure 1: Process of localizing task-sensitive attention heads. The left side represents the attention heads of LLM, and the center illustrates the distribution of head parameters, where task parameters \( \theta_T \) capture task information and lead to task-sensitive heads \( H_T \).
Figure 2: Preliminary results comparing ALPS against baselines across different LLM scales (Left), and downstream tasks (Right). ALPS consistently outperforms other methods, demonstrating its efficiency in diverse settings.

look the intrinsic functionality of the model weight parameters, which inherently encode task-relevant information (Zhong et al., 2022; Tam et al., 2024), and can be leveraged to identify task-sensitive attention heads without relying on activation data.

In this work, to enhance alignment efficiency, utilize weight parameters, and introduce reusability of identified attention heads, we propose the Attention Localization and Pruning Strategy (ALPS), a heuristic-guided algorithm that localizes task-sensitive attention heads and prunes by restricting attention training updates to these heads only. This algorithm draws inspiration from prior work leveraging model weight parameters to localize task-sensitive attention heads (Voita et al., 2019; Zheng et al., 2024; Shi et al., 2024), as shown in Fig 1. Unlike previous work utilizing whole model parameters, ALPS focuses on the distribution of each attention head. Specifically, given a base model and its task-related model, we extract and normalize the distribution of each head, then quantify the distributional shift between corresponding heads using our proposed parameter alignment distribution score. The heads with the highest scores are then selected as task-sensitive heads, while the remaining heads are frozen during fine-tuning. This yields a pruned, task-specialized model with enhanced efficiency. Empirical results demonstrate that our method updates only \( 10\% \) of attention parameters during fine-tuning while achieving a \( 2\% \) performance gain over baselines across general, math, and code tasks, as shown in Figure 2. Further analysis reveals that the selected task-sensitive heads exhibit transferability and reusability, enhancing performance across datasets within the same task, reducing computational overhead, and avoiding data dependency. Moreover, the induced sparsity in attention heads helps mitigate knowledge forgetting in the model alignment.

In summary, our contributions are as follows:

• We offer a novel perspective on efficient LLM alignment, pioneering an approach using model weight parameters to identify task-sensitive attention heads.

• We propose ALPS, a novel algorithm that identifies task-sensitive attention heads by measuring shifts in their weight distributions and then restricts training to these heads.

• Empirical results demonstrate that our method improves alignment efficiency and task performance while introducing sparsity and head transferability.

2 Related Work

2.1 Efficient LLMs
To mitigate the high costs of aligning foundation LLMs to downstream tasks, researchers have explored various strategies for efficient alignment (Wan et al., 2023). For example, DeepSeek-R1 (Guo et al., 2025) relies solely on reinforcement training to reduce overhead. While data-focused approaches either aim to minimize the need for task-specific data (Zhou et al., 2024a; Chen et al., 2023a,b) or augment downstream datasets (Li et al., 2023b; Wang et al., 2022; Xu et al., 2023), parameter-focused approaches improve efficiency by avoiding full-model updates (Hu et al., 2021), compressing the model (Jacob et al., 2018; Hinton, 2015), and streamlining key-value storage (Ainslie et al., 2023; Liu et al., 2024). Our work leverages the attention mechanism to localize and prune task-sensitive attention heads.

2.2 Attention Localization and Pruning
Recent studies have investigated attention mechanisms at various granularities to improve the efficiency and interpretability of LLM, including neurons, layers, and heads (Geiger et al., 2021; Gurnee et al., 2023; Zou et al., 2023; Zheng et al., 2024; Zhao et al., 2024). Notably, many works have concentrated on attention heads for downstream
Figure 3: Overview of ALPS framework. a) Given a base model \( \mathcal{M}_B \) and a task fine-tuned model \( \mathcal{M}_T \), we extract their attention heads, compute weight distributions using softmax, and calculate their PAD score \( s_h^{PAD} \). The Top-K attention heads are then selected based on \( S^{PAD} \), while b) pruning the remaining heads by freezing their gradient updates during fine-tuning. The resulting model \( \mathcal{M}_{ALPS} \) outperforms \( \mathcal{M}_T \) in both performance and efficiency, with transferable task-relevant heads that enhance alignment across datasets within the same task.

tasks, including general linguistic and factuality capability (Clark, 2019; Wu et al., 2024; Zheng et al., 2024), model safety (Chen et al., 2024; Zhou et al., 2024b), KV-cache compression (Ainslie et al., 2023; Fu et al., 2024), etc. Among them, Michel et al. (2019) demonstrated that most attention heads can be dropped without compromising performance. While most methods rely on the activation of model parameters by feeding task-specific data, Voita et al. (2019) analyzed the weight matrices of attention models, revealing that specialized heads contribute to tasks, while the rest can be pruned. Similarly, Shi et al. (2024) and He et al. (2024) also observed that pruning certain attention layers enhances alignment efficiency, indicating redundancy within the attention mechanism. In contrast, our approach directly localizes task-relevant attention heads from weight matrices (Voita et al., 2019; Shi et al., 2024), thereby reducing data dependency and improving generalizability and efficiency in downstream alignment.

3 Methodology

In this section, we first provide a preliminary overview of the attention mechanism with task parameters and then present ALPS in detail. Fig 3 illustrates the overall framework.

3.1 Preliminary

In this work, we adopt the Llama-3 series of models, which employ grouped-query attention (GQA) (Ainslie et al., 2023) to reduce KV cache overhead. Formally, given an input sequence matrix \( X \in \mathbb{R}^{t \times d} \), where \( t \) represents the sequence length, GQA divides the \( n \) attention heads into \( g \) groups, making all heads within the same group share a single pair of KV projections. Each head \( h \in \{1, \ldots, n\} \) learns projections \( \{W_q^h, W_k^{[hg/n]}, W_v^{[hg/n]}\} \), where \( \lceil \cdot \rceil \) is the ceiling function, and mapping \( X \) into queries \( Q^h \), keys \( K^{[hg/n]} \), and values \( V^{[hg/n]} \):

\[
Q^h = X W_q^h \in \mathbb{R}^{t \times d_k}, \\
K^{[hg/n]} = X W_k^{[hg/n]} \in \mathbb{R}^{t \times d_k}, \\
V^{[hg/n]} = X W_v^{[hg/n]} \in \mathbb{R}^{t \times d_v},
\]

where a common choice sets \( d_k = d_v = d_{model}/h \), and if \( g = h \), this GQA setting reduces to standard multi-head attention. In the Llama-3 collection, a typical choice sets \( g = 8 \) for KV sharing. The output of head \( s \) is then computed as:

\[
O^h = \text{Softmax}\left( \frac{Q^h (K^{[hg/n]})^\top}{\sqrt{d_k}} \right) V^{[hg/n]}, \\
O^h \in \mathbb{R}^{t \times d_v}.
\]

3.2 Task Head Parameters

Downstream task-related head parameters are critical components in LLMs that directly influence performance on specific tasks. Inspired by mechanistic interpretability studies (Voita et al., 2019; Zhao et al., 2024; Lindner et al., 2023), we define these task head parameters as those whose ablation results in a significant degradation of task-specific
performance. Formally, given a downstream task \( \mathcal{T} \) with evaluation metric \( p \), we quantify the impact of ablating a head parameter \( \theta^h \) by the relative performance drop:

\[
\Delta p(\theta^h) = p\left(\theta^{\mathcal{M}}, \mathcal{T}\right) - p\left(\theta^{\mathcal{M}} \setminus \theta^h, \mathcal{T}\right),
\]

where \( \theta^{\mathcal{M}} \) denotes the original model parameters, and \( \theta^{\mathcal{M}} \setminus \theta^h \) represents the model after ablating head parameters \( \theta^h \). The downstream task-related head parameters are then identified as the Top-K heads that maximize this performance drop:

\[
\Theta_{\mathcal{T}, K} = \operatorname{Top-K} \left\{ \theta^{\mathcal{T}} : \underset{\theta^h \in \theta^{\mathcal{M}}}{\operatorname{argmax}} \Delta p(\theta^h) \right\}.
\]

The ablation of \( \Theta_{\mathcal{T}, K} \) leads to a measurable decline in task performance, underscoring their importance.

3.3 Attention Localization and Pruning Strategy (ALPS)

To localize task-sensitive attention heads that contribute to downstream task alignment, we introduce a heuristic search algorithm named \emph{Attention Localization and Pruning Strategy (ALPS)}. The algorithm is divided into two stages: \textbf{Localizing} task-sensitive attention heads and \textbf{Pruning} non-critical heads in task alignment.

\textbf{Localizing.} To identify task-sensitive attention heads, we introduce the \emph{parameter alignment distribution score} \( s^{PAD} \), which is based on the Wasserstein-1 distance (Vaserstein, 1969) between the weight matrices of a fine-tuned task model and its base counterpart, quantifying the variation in attention head parameters across different tasks.

Given a pre-trained model \( \mathcal{M}_B \) and its task fine-tuned version \( \mathcal{M}_T \), we analyze the static projection matrices of each attention head \( h \). For head \( h \) with parameters \( \{ W_q^h, W_k^{[hg/n]}, W_v^{[hg/n]} \} \), the projection matrix is computed as:

\[
W_o^h = W_q^h W_k^{[hg/n]^\top} W_v^{[hg/n]} \in \mathbb{R}^{d \times d},
\]

which captures the transformation behavior of one head(Kobayashi et al., 2020).

To measure task-induced parameter shifts, we first convert \( W_o^{(s)} \) into a probability distribution via tempered softmax:

\[
P^h = \operatorname{Softmax}\left( \frac{W_o^h}{\tau} \right), \quad \tau > 0,
\]

where \( \tau \) controls distribution granularity, and we set \( \tau \) to 1 for simplicity. The Wasserstein-1 distance between base and task distributions then quantifies head sensitivity:

\[
s_h^{PAD} = W_1 \left( P_B^h, P_T^h \right).
= \inf_{\gamma \in \Gamma(P_B^h, P_T^h)} \mathbb{E}_{(x, y) \sim \gamma} \left[ \| x - y \| \right],
\]

where \( \Gamma \) denotes all joint distributions with marginals \( P_B^h \) and \( P_T^h \), and we compute \( s_h^{PAD} \) for each attention head. Results in Section 4.3 show that heads with higher \( s_h^{PAD} \) consistently achieve stronger task alignment and thus better performance, confirming that the magnitude of parameter variation reflects their importance for downstream tasks.

\textbf{Pruning.} After obtaining a set of \( S^{PAD} = \{ s_h^{PAD} | h \in \mathcal{H} \} \), where \( \mathcal{H} = \{ 1, \ldots, n \} \) denote all attention heads in the model, we select a group of task-sensitive heads \( \mathcal{H}_T \) with Top-K\( (S^{PAD}) \) to retain. During task fine-tuning, the remaining heads will be pruned by freezing gradients of these heads, which masks their parameter updates to 0. Specifically, for all attention heads, only parameters associated with \( \mathcal{H}_T \) are updated:

\[
\nabla_{\theta^h} \mathcal{L} = \begin{cases}
\nabla_{\theta^h} \mathcal{L} & \text{if } h \in \mathcal{H}_S \\
0 & \text{otherwise}
\end{cases}
\]

Algorithm 1 Attention Localization and Pruning Strategy (ALPS)

Input: \( \mathcal{M}_B, \mathcal{M}_T, \mathcal{D}_T \), retention ratio \( r \in (0, 1] \)
Output: \( \mathcal{H}_r, \mathcal{M}_{ALPS} \)
1: Let \( \mathcal{H} \gets \{ 1, 2, \ldots, n \} \triangleright \) Set of all attention heads
2: \textbf{for all} \( h \in \mathcal{H} \) \textbf{do}
3: \hspace{0.5em} \( P_B^h \gets \operatorname{Softmax}(W_{o,B}^h) \)
4: \hspace{0.5em} \( P_T^h \gets \operatorname{Softmax}(W_{o,T}^h) \)
5: \hspace{0.5em} Compute \( s_h^{PAD} \gets W_1(P_B^h, P_T^h) \)
6: \textbf{end for}
7: \( K \gets \lceil r \cdot n \rceil \)
8: \( S^{PAD} \gets \{ s_h^{PAD} | h \in \mathcal{H} \} \)
9: \( \mathcal{H}_T \gets \operatorname{Top-K}(S^{PAD}) \triangleright \) the set of heads with Top-K \( s_h^{PAD} \)
10: \textbf{for all} \( \hat{h} \in \mathcal{H} \setminus \mathcal{H}_T \) \textbf{do}
11: \hspace{0.5em} \( \nabla_{\theta^{\hat{h}}} \mathcal{L} = 0 \triangleright \) Ablate non-critical heads during fine-tuning
12: \textbf{end for}
13: Fine-tune model on \( \mathcal{D}_T \) to obtain \( \mathcal{M}_{ALPS} \)
14: \textbf{return} \( \mathcal{H}_r, \mathcal{M}_{ALPS} \)
<table>
  <tr>
    <th rowspan="2">Model</th>
    <th rowspan="2">Method</th>
    <th rowspan="2">Attn%</th>
    <th colspan="3">General</th>
    <th colspan="3">Math</th>
    <th colspan="3">Code</th>
    <th rowspan="2">Avg.</th>
  </tr>
  <tr>
    <th>IEval</th>
    <th>GPQA</th>
    <th>GSM8K</th>
    <th>MATH</th>
    <th>HEval</th>
    <th>HEval+</th>
    <th>MBPP</th>
    <th>MBPP+</th>
  </tr>
  <tr>
    <td rowspan="5">Llama-3.2-1B</td>
    <td>w/ full</td>
    <td>100%</td>
    <td>34.53</td>
    <td>18.75</td>
    <td>17.82</td>
    <td>4.60</td>
    <td>45.12</td>
    <td>39.08</td>
    <td>28.63</td>
    <td>23.78</td>
    <td>26.54</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>0%</td>
    <td>35.85</td>
    <td>19.64</td>
    <td>23.35</td>
    <td>5.20</td>
    <td>41.51</td>
    <td>35.28</td>
    <td>29.91</td>
    <td><b>24.88</b></td>
    <td>26.95 ±0.41</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>10%</td>
    <td>33.29</td>
    <td>25.89</td>
    <td>22.82</td>
    <td>5.02</td>
    <td>41.51</td>
    <td>37.22</td>
    <td>28.12</td>
    <td>24.13</td>
    <td>27.25 ±0.71</td>
  </tr>
  <tr>
    <td>LC</td>
    <td>10%</td>
    <td>32.15</td>
    <td>27.88</td>
    <td>21.73</td>
    <td>4.81</td>
    <td>42.05</td>
    <td>38.07</td>
    <td>27.03</td>
    <td>23.88</td>
    <td>27.20 ±0.66</td>
  </tr>
  <tr>
    <td>LoRA</td>
    <td>~10%</td>
    <td><b>36.75</b></td>
    <td>26.14</td>
    <td>23.67</td>
    <td>5.92</td>
    <td>42.34</td>
    <td>37.99</td>
    <td>28.92</td>
    <td>24.10</td>
    <td>28.22 ±1.68</td>
  </tr>
  <tr>
    <td rowspan="5">Llama-3.2-3B</td>
    <td>w/ full</td>
    <td>100%</td>
    <td>42.81</td>
    <td>22.32</td>
    <td>31.77</td>
    <td>16.68</td>
    <td>56.18</td>
    <td>50.53</td>
    <td>50.00</td>
    <td>41.50</td>
    <td>38.97</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>0%</td>
    <td>40.29</td>
    <td>24.33</td>
    <td>33.21</td>
    <td>13.76</td>
    <td>55.52</td>
    <td>50.62</td>
    <td>48.71</td>
    <td>39.72</td>
    <td>38.27 ±0.70</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>10%</td>
    <td>42.21</td>
    <td>20.54</td>
    <td>35.03</td>
    <td>14.56</td>
    <td>56.18</td>
    <td>49.43</td>
    <td>54.57</td>
    <td>45.28</td>
    <td>39.73 ±0.76</td>
  </tr>
  <tr>
    <td>LC</td>
    <td>10%</td>
    <td>41.75</td>
    <td>21.86</td>
    <td><b>36.42</b></td>
    <td>13.94</td>
    <td>55.82</td>
    <td>50.67</td>
    <td>55.13</td>
    <td>42.50</td>
    <td>39.76 ±0.79</td>
  </tr>
  <tr>
    <td>LoRA</td>
    <td>~10%</td>
    <td>42.90</td>
    <td>21.22</td>
    <td>34.20</td>
    <td>15.36</td>
    <td>56.85</td>
    <td>50.09</td>
    <td><b>55.27</b></td>
    <td><b>45.03</b></td>
    <td>40.12 ±1.15</td>
  </tr>
  <tr>
    <td rowspan="5">Llama-3.1-8B</td>
    <td>w/ full</td>
    <td>100%</td>
    <td>51.20</td>
    <td>25.22</td>
    <td>63.00</td>
    <td>20.52</td>
    <td>69.50</td>
    <td>64.28</td>
    <td>63.21</td>
    <td>52.38</td>
    <td>51.16</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>0%</td>
    <td>40.41</td>
    <td>26.16</td>
    <td>64.06</td>
    <td>15.48</td>
    <td>70.72</td>
    <td><b>65.88</b></td>
    <td>63.27</td>
    <td>52.55</td>
    <td>49.82 ±1.34</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>10%</td>
    <td>49.52</td>
    <td>26.34</td>
    <td>62.40</td>
    <td>23.52</td>
    <td>70.12</td>
    <td>63.39</td>
    <td>59.00</td>
    <td>49.21</td>
    <td>50.44 ±0.72</td>
  </tr>
  <tr>
    <td>LC</td>
    <td>10%</td>
    <td>48.67</td>
    <td>27.10</td>
    <td>63.83</td>
    <td>22.47</td>
    <td>70.68</td>
    <td>63.05</td>
    <td>58.28</td>
    <td>49.79</td>
    <td>50.48 ±0.68</td>
  </tr>
  <tr>
    <td>LoRA</td>
    <td>~10%</td>
    <td>50.02</td>
    <td>26.82</td>
    <td>62.92</td>
    <td><b>24.07</b></td>
    <td>72.50</td>
    <td>63.89</td>
    <td>59.55</td>
    <td>49.69</td>
    <td>51.18 ±0.02</td>
  </tr>
  <tr>
    <td rowspan="1">ALPS</td>
    <td>10%</td>
    <td><b>51.33</b></td>
    <td><b>27.29</b></td>
    <td><b>64.13</b></td>
    <td>22.48</td>
    <td><b>72.68</b></td>
    <td>65.03</td>
    <td><b>63.67</b></td>
    <td><b>52.88</b></td>
    <td><b>52.41</b> ±1.25</td>
  </tr>
</table>

Table 1: Performance of our method compared to other baselines with Llama-3.2-1B, Llama-3.2-3B, and Llama-3.1-8B in general, math, and code tasks. The best results are highlighted in bold. ↑ indicates change relative to the w/ full baseline. Note that each task uses a separately fine-tuned model trained on its respective dataset. ALPS achieves better performance while updating only **10%** of attention parameters, highlighting its efficiency.

where \( \theta^h = \{ W_q^{(h)}, W_k^{(h)}, W_v^{(h)} \} \). This gradient masking ensures that non-sensitive heads \( \mathcal{H} \setminus \mathcal{H}_T \) retain pre-trained knowledge without interference, and the loss objective turns to:

\[
\mathcal{L} = \mathbb{E}_{(x,y) \sim \mathcal{D}_T} \left[ -\log p \left( y \mid x; \{ \theta^h \}_{h \in \mathcal{H}_T} \right) \right],
\]

where \( \{ \theta_s \} \) denotes the parameters of heads in \( \mathcal{H}_S \), and \( \mathcal{D}_T \) represents the downstream task data. This approach reduces optimization redundancy by focusing updates on heads critical to \( \mathcal{T} \).

The overall procedure is detailed in Algorithm 1, and the final output of the algorithm is the set of sensitive attention heads \( \mathcal{H}_T \) with modified fine-tuned model \( \mathcal{M}_{ALPS} \).

4 Experiments

4.1 Setup

Datasets. We evaluate our approach on three representative datasets covering general, math, and code tasks: UltraChat (Ding et al., 2023), MathInstruct (Yue et al., 2023), and MagicCoder (Wei et al., 2023). To assess the transferability of our method, we further conduct experiments on Alpaca (Taori et al., 2023), Camel-math (Li et al., 2023a), and CodeAlpaca (Chaudhary, 2023). Details are listed in Appendix A.

Models. We conduct our method on the Llama-3 series models (Llama-3.2-1B, Llama-3.2-3B, and Llama-3.18B), all pre-trained with GQA for improved efficiency. Each model shares key and value projections across 8 heads (\( q=8 \)), while queries remain head-specific. Thus, when computing \( s^{PAD} \), each key-value pair corresponds to 8 query heads.

Evaluation. We use lm-eval (Gao et al., 2024; meta llama, 2023) to evaluate general capabilities in instruction following, reasoning, and dialogue, as well as mathematical abilities. For code generation, we employ EvalPlus (Liu et al., 2023). All evaluations follow the official implementations, with details on benchmarks and metrics provided in Appendix A.

Baseline. To evaluate the effectiveness of ALPS , we compare it against the following baselines. **w/ full**: standard full supervised fine-tuning where all attention parameters are updated. **w/o attn**: freeze all attention heads by preventing updates to the QKV weight matrices. **Random**: selects a fixed proportion of random attention heads for training. **LC**: layer consistency selection, extending random by ensuring a uniform distribution of selected heads across layers, maintaining fairness in head allocation. **LoRA** (Hu et al., 2021): a
<table>
  <tr>
    <th rowspan="2">Model</th>
    <th rowspan="2">Method</th>
    <th rowspan="2">Attn%</th>
    <th colspan="2">General</th>
    <th colspan="2">Math</th>
    <th colspan="2">Code</th>
    <th rowspan="2">Avg.</th>
  </tr>
  <tr>
    <th>IFEval</th>
    <th>GPQA</th>
    <th>GSM8K</th>
    <th>MATH</th>
    <th>HEval</th>
    <th>HEval+</th>
    <th>MBPP</th>
    <th>MBPP+</th>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.2-1B</td>
    <td>C.S.</td>
    <td>10%</td>
    <td>28.12</td>
    <td>17.97</td>
    <td>18.18</td>
    <td>4.15</td>
    <td>28.22</td>
    <td>23.97</td>
    <td>18.89</td>
    <td>13.67</td>
    <td>19.15</td>
  </tr>
  <tr>
    <td>Eu.</td>
    <td>10%</td>
    <td>34.05</td>
    <td>24.68</td>
    <td>21.91</td>
    <td>6.89</td>
    <td>44.55</td>
    <td>39.12</td>
    <td>27.14</td>
    <td>22.54</td>
    <td>27.61</td>
  </tr>
  <tr>
    <td>KL</td>
    <td>10%</td>
    <td>35.22</td>
    <td><b>26.45</b></td>
    <td><b>23.95</b></td>
    <td>7.05</td>
    <td>45.83</td>
    <td><b>40.45</b></td>
    <td>28.72</td>
    <td><b>24.55</b></td>
    <td>29.03</td>
  </tr>
  <tr>
    <td>s<sup>SPAD</sup></td>
    <td>10%</td>
    <td><b>36.69</b></td>
    <td>26.16</td>
    <td>23.74</td>
    <td><b>7.28</b></td>
    <td><b>46.89</b></td>
    <td>40.22</td>
    <td><b>29.13</b></td>
    <td>24.21</td>
    <td><b>29.29</b></td>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.2-3B</td>
    <td>C.S.</td>
    <td>10%</td>
    <td>32.57</td>
    <td>19.88</td>
    <td>22.05</td>
    <td>10.67</td>
    <td>46.12</td>
    <td>38.83</td>
    <td>38.97</td>
    <td>30.25</td>
    <td>29.92</td>
  </tr>
  <tr>
    <td>Eu.</td>
    <td>10%</td>
    <td>42.36</td>
    <td>22.45</td>
    <td>32.11</td>
    <td>16.23</td>
    <td>55.67</td>
    <td>49.25</td>
    <td>49.82</td>
    <td>40.15</td>
    <td>38.51</td>
  </tr>
  <tr>
    <td>KL</td>
    <td>10%</td>
    <td>43.78</td>
    <td><b>24.50</b></td>
    <td>33.14</td>
    <td>16.78</td>
    <td>56.89</td>
    <td>50.33</td>
    <td>50.47</td>
    <td>41.12</td>
    <td>39.63</td>
  </tr>
  <tr>
    <td>s<sup>SPAD</sup></td>
    <td>10%</td>
    <td><b>44.96</b></td>
    <td>24.33</td>
    <td><b>34.27</b></td>
    <td><b>17.12</b></td>
    <td><b>58.28</b></td>
    <td><b>51.67</b></td>
    <td><b>52.21</b></td>
    <td><b>42.78</b></td>
    <td><b>40.70</b></td>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.1-8B</td>
    <td>C.S.</td>
    <td>10%</td>
    <td>41.85</td>
    <td>22.47</td>
    <td>48.22</td>
    <td>19.17</td>
    <td>61.45</td>
    <td>52.93</td>
    <td>55.34</td>
    <td>49.05</td>
    <td>43.81</td>
  </tr>
  <tr>
    <td>Eu.</td>
    <td>10%</td>
    <td>48.94</td>
    <td>26.03</td>
    <td>62.78</td>
    <td>20.84</td>
    <td>69.12</td>
    <td>63.47</td>
    <td>61.02</td>
    <td>50.37</td>
    <td>50.32</td>
  </tr>
  <tr>
    <td>KL</td>
    <td>10%</td>
    <td>50.12</td>
    <td>26.75</td>
    <td>63.91</td>
    <td>21.26</td>
    <td><b>73.05</b></td>
    <td>64.12</td>
    <td>62.15</td>
    <td>51.43</td>
    <td>51.60</td>
  </tr>
  <tr>
    <td>s<sup>SPAD</sup></td>
    <td>10%</td>
    <td><b>51.33</b></td>
    <td><b>27.13</b></td>
    <td><b>64.13</b></td>
    <td><b>22.48</b></td>
    <td>72.68</td>
    <td><b>65.03</b></td>
    <td><b>63.67</b></td>
    <td><b>52.88</b></td>
    <td><b>52.41</b></td>
  </tr>
</table>

Table 2: Ablation study on metrics for ALPS with 10% of attention heads. Bold indicates the best performance.

<table>
  <tr>
    <th rowspan="2">Model</th>
    <th rowspan="2">Ratio</th>
    <th colspan="2">General</th>
    <th colspan="2">Math</th>
    <th colspan="2">Code</th>
    <th rowspan="2">Avg.</th>
  </tr>
  <tr>
    <th>IFEval</th>
    <th>GPQA</th>
    <th>GSM8K</th>
    <th>MATH</th>
    <th>HEval</th>
    <th>HEval+</th>
    <th>MBPP</th>
    <th>MBPP+</th>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.2-1B</td>
    <td>100%</td>
    <td>34.53</td>
    <td>18.75</td>
    <td>17.82</td>
    <td>4.60</td>
    <td>45.12</td>
    <td>39.08</td>
    <td>28.63</td>
    <td>23.78</td>
    <td>26.54</td>
  </tr>
  <tr>
    <td>70%</td>
    <td>31.80</td>
    <td>16.91</td>
    <td>20.63</td>
    <td>6.35</td>
    <td>44.20</td>
    <td>39.50</td>
    <td>28.40</td>
    <td>24.90</td>
    <td>26.59 <span style="color:blue;">[0.05]</span></td>
  </tr>
  <tr>
    <td>50%</td>
    <td>33.46</td>
    <td>18.07</td>
    <td>21.08</td>
    <td>7.15</td>
    <td>43.95</td>
    <td>39.69</td>
    <td>28.92</td>
    <td>24.64</td>
    <td>27.12 <span style="color:blue;">[0.58]</span></td>
  </tr>
  <tr>
    <td>30%</td>
    <td><b>37.12</b></td>
    <td>25.98</td>
    <td>22.56</td>
    <td>7.12</td>
    <td>46.34</td>
    <td>39.84</td>
    <td>28.95</td>
    <td>24.05</td>
    <td>29.00 <span style="color:blue;">[2.46]</span></td>
  </tr>
  <tr>
    <td>10%</td>
    <td><b>36.69</b></td>
    <td><b>26.16</b></td>
    <td><b>23.74</b></td>
    <td><b>7.28</b></td>
    <td><b>46.89</b></td>
    <td><b>40.22</b></td>
    <td>29.13</td>
    <td>24.21</td>
    <td><b>29.29</b> <span style="color:blue;">[2.75]</span></td>
  </tr>
  <tr>
    <td>0%</td>
    <td>35.85</td>
    <td>19.64</td>
    <td>23.35</td>
    <td>5.20</td>
    <td>41.51</td>
    <td>35.28</td>
    <td><b>29.91</b></td>
    <td><b>24.88</b></td>
    <td>26.95 <span style="color:blue;">[0.41]</span></td>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.2-3B</td>
    <td>100%</td>
    <td>42.81</td>
    <td>22.32</td>
    <td>31.77</td>
    <td>16.68</td>
    <td>56.18</td>
    <td>50.53</td>
    <td>50.00</td>
    <td>41.50</td>
    <td>38.97</td>
  </tr>
  <tr>
    <td>70%</td>
    <td>39.83</td>
    <td>21.91</td>
    <td>31.63</td>
    <td>15.35</td>
    <td>56.21</td>
    <td>51.55</td>
    <td>49.40</td>
    <td>39.97</td>
    <td>38.23 <span style="color:blue;">[0.74]</span></td>
  </tr>
  <tr>
    <td>50%</td>
    <td>40.93</td>
    <td>22.18</td>
    <td>32.17</td>
    <td>16.10</td>
    <td>55.91</td>
    <td>49.72</td>
    <td>50.05</td>
    <td>40.33</td>
    <td>38.42 <span style="color:blue;">[0.55]</span></td>
  </tr>
  <tr>
    <td>30%</td>
    <td>43.53</td>
    <td>23.46</td>
    <td>33.91</td>
    <td>16.54</td>
    <td>57.28</td>
    <td>51.69</td>
    <td>51.61</td>
    <td>41.93</td>
    <td>39.99 <span style="color:blue;">[1.02]</span></td>
  </tr>
  <tr>
    <td>10%</td>
    <td><b>44.96</b></td>
    <td><b>24.33</b></td>
    <td><b>34.27</b></td>
    <td><b>17.12</b></td>
    <td><b>58.28</b></td>
    <td><b>51.67</b></td>
    <td><b>52.21</b></td>
    <td><b>42.78</b></td>
    <td><b>40.70</b> <span style="color:blue;">[1.73]</span></td>
  </tr>
  <tr>
    <td>0%</td>
    <td>40.29</td>
    <td>24.33</td>
    <td>33.21</td>
    <td>13.76</td>
    <td>55.52</td>
    <td>50.62</td>
    <td>48.71</td>
    <td>39.72</td>
    <td>38.27 <span style="color:blue;">[0.70]</span></td>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.1-8B</td>
    <td>100%</td>
    <td>51.20</td>
    <td>25.22</td>
    <td>63.00</td>
    <td>20.52</td>
    <td>69.50</td>
    <td>64.28</td>
    <td>63.21</td>
    <td>52.38</td>
    <td>51.16</td>
  </tr>
  <tr>
    <td>70%</td>
    <td>50.80</td>
    <td>24.91</td>
    <td>62.63</td>
    <td>20.35</td>
    <td>69.20</td>
    <td>61.50</td>
    <td>63.40</td>
    <td>51.90</td>
    <td>50.59 <span style="color:blue;">[0.57]</span></td>
  </tr>
  <tr>
    <td>50%</td>
    <td><b>52.55</b></td>
    <td><b>27.28</b></td>
    <td>61.11</td>
    <td>20.95</td>
    <td>69.69</td>
    <td>63.35</td>
    <td>61.71</td>
    <td>51.02</td>
    <td>50.96 <span style="color:blue;">[0.20]</span></td>
  </tr>
  <tr>
    <td>30%</td>
    <td>51.85</td>
    <td>26.47</td>
    <td>62.72</td>
    <td>22.15</td>
    <td>71.15</td>
    <td>64.55</td>
    <td>62.88</td>
    <td>52.23</td>
    <td>51.62 <span style="color:blue;">[0.46]</span></td>
  </tr>
  <tr>
    <td>10%</td>
    <td><b>51.33</b></td>
    <td><b>27.13</b></td>
    <td><b>64.13</b></td>
    <td><b>22.48</b></td>
    <td><b>72.68</b></td>
    <td><b>65.03</b></td>
    <td><b>63.67</b></td>
    <td><b>52.88</b></td>
    <td><b>52.41</b> <span style="color:blue;">[1.25]</span></td>
  </tr>
  <tr>
    <td>0%</td>
    <td>40.41</td>
    <td>26.16</td>
    <td>64.06</td>
    <td>15.48</td>
    <td>70.72</td>
    <td>64.88</td>
    <td>63.27</td>
    <td>52.55</td>
    <td>49.69 <span style="color:blue;">[1.47]</span></td>
  </tr>
</table>

Table 3: Ablation study on attention head selection ratio for ALPS . The best results are highlighted in bold. ↓↑ indicates change relative to the w/ full baseline.

widely used parameter-efficient fine-tuning method that freezes all original weights and updates only a low-rank decomposition of weight matrices. We set the LoRA rank to 128 and only to QKV matrices. For consistency in head indexing, all head selections follow the query index. For a Q index of 20, the K and V indices are determined by \( [20 * g / h] \), where \( g \) is the number of KV groups, and \( h \) is the total number of heads.

Implementation Details. We optimize our models using the AdamW (Loshchilov, 2017) optimizer with hyperparameters \( \beta_1 = 0.9 \) and \( \beta_2 = 0.999 \), and a weight decay of 0.1. A cosine decay schedule is employed, gradually reducing the learning rate to 10% of its initial value throughout training, following a linear warmup ratio of 0.1. We set the effective batch size to 128 and initialize the learning rate at \( 2 \times 10^{-5} \). All models are trained for 3 epochs using mixed precision (fp16) on 8 A100 GPUs.

4.2 Main Results

Table 1 presents the main results comparing ALPS against other baselines across three model scales (1B, 3B, and 8B) and three downstream tasks (general, math, and code). Note that each task uses a separately fine-tuned model trained on its respec-
tive dataset, that is, the results of general tasks are from the general fine-tuned model, while the results of math and code are from math and code fine-tuned models, respectively. Notably, our method consistently outperforms all baselines, showing a 2.75% improvement over the w/ full baseline on the 1B model and maintains a 1.25% gain at the 8B scale. This indicates the redundancy among attention heads in downstream alignment. In contrast, the w/o attn settings fall significantly behind the w/ full baseline, highlighting the critical role of attention heads in downstream alignment. Furthermore, results from Random, LC, LoRA, and our method indicate that selecting a subset of heads can indeed enhance alignment efficiency. However, how these heads are chosen remains pivotal, and selecting more task-sensitive attention heads can outperform training on all heads, underscoring the importance of precise head identification for maximizing both efficiency and performance.

![Ablation study on attention head selection ratio. The efficiency ratio indicates that activating 10% and 30% of attention heads yields better performance than full fine-tuning across three models.](page_146_693_579_312.png)

Figure 4: Ablation study on attention head selection ratio. The efficiency ratio indicates that activating 10% and 30% of attention heads yields better performance than full fine-tuning across three models.

**4.3 Ablations Study**

**Head Metric Ablation.** Table 2 compares four metrics: C.S. (cosine similarity), Eu. (Euclidean distance), KL (Kullback-Leibler Divergence (Kullback and Leibler, 1951)), and our proposed parameter alignment distribution score \( s^{PAD} \), for identifying task-sensitive attention heads. Notably, C.S. yields the weakest performance, likely because it focuses on overall matrix similarity, which can overlook significant local shifts tied to specific task requirements. In contrast, across all model scales (1B, 3B, 8B) and tasks (general, math, code), our proposed metric \( s^{PAD} \) consistently achieves the highest average scores, indicating its effectiveness in capturing the most critical distribution shifts between the base and task fine-tuned models.

**Head Ratio Ablation.** The ablation study on the proportion of attention heads selected for our method ranging from 0% to 100%, is presented in Table 3 and Figure 4. We observe that both 10% and 30% ratios consistently yield strong performance across all model scales and tasks, as highlighted by the efficiency ratio region in Figure 4. Notably, updating 10% of the heads already achieves comparable to or better results than higher ratios, and strikes a better trade-off between performance and computational cost by reducing the number of trainable parameters. Consequently, we adopt 10% as our default setting for improved efficiency.

![Heatmaps of task-sensitive attention heads (10%) selected by our method on Llama-3.2-1B for a) general, b) math, and c) code tasks. The color intensity indicates the value of \( s^{PAD} \) of each head.](page_820_1042_579_312.png)

Figure 5: Heatmaps of task-sensitive attention heads (10%) selected by our method on Llama-3.2-1B for a) general, b) math, and c) code tasks. The color intensity indicates the value of \( s^{PAD} \) of each head.

**4.4 Head Impact Analysis**

Figure 5 shows heatmaps of the top 10% attention heads identified by our method across three tasks, full heatmaps are listed in Appendix F. Each grid cell corresponds to a layer-head combination, and the color intensity reflects the importance of the head. We observe that general and code tasks exhibit notable overlap in their most sensitive heads, potentially due to the presence of extensive natural language instructions in the code datasets. Meanwhile, math reveals a more distinct pattern, suggest-
<table>
  <tr>
    <th>Model</th>
    <th>Method</th>
    <th>General</th>
    <th>Math</th>
    <th>Code</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.2-1B</td>
    <td>w/ full</td>
    <td>14.28</td>
    <td>5.87</td>
    <td>22.28</td>
    <td>14.14</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>19.73</td>
    <td>8.39</td>
    <td>19.98</td>
    <td>16.03 <span style="color:red;">†1.89</span></td>
  </tr>
  <tr>
    <td>Random</td>
    <td>21.80</td>
    <td>7.91</td>
    <td>19.63</td>
    <td>16.45 <span style="color:red;">†2.31</span></td>
  </tr>
  <tr>
    <td><b>ALPS</b></td>
    <td><b>23.43</b></td>
    <td><b>11.51</b></td>
    <td><b>24.11</b></td>
    <td><b>19.68 <span style="color:red;">†5.54</span></b></td>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.2-3B</td>
    <td>w/ full</td>
    <td>21.57</td>
    <td>15.23</td>
    <td>38.55</td>
    <td>25.12</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>22.31</td>
    <td>13.49</td>
    <td>35.72</td>
    <td>23.84 <span style="color:red;">†1.28</span></td>
  </tr>
  <tr>
    <td>Random</td>
    <td>22.38</td>
    <td>17.80</td>
    <td>41.37</td>
    <td>27.18 <span style="color:red;">†2.06</span></td>
  </tr>
  <tr>
    <td><b>ALPS</b></td>
    <td><b>24.65</b></td>
    <td><b>18.70</b></td>
    <td><b>40.24</b></td>
    <td><b>27.86 <span style="color:red;">†2.74</span></b></td>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.1-8B</td>
    <td>w/ full</td>
    <td>27.21</td>
    <td>28.76</td>
    <td>42.34</td>
    <td>32.77</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>18.29</td>
    <td>25.77</td>
    <td>46.11</td>
    <td>30.06 <span style="color:red;">†3.38</span></td>
  </tr>
  <tr>
    <td>Random</td>
    <td>27.93</td>
    <td>32.96</td>
    <td>40.43</td>
    <td>33.77 <span style="color:red;">†1.00</span></td>
  </tr>
  <tr>
    <td><b>ALPS</b></td>
    <td><b>28.21</b></td>
    <td><b>31.31</b></td>
    <td><b>44.57</b></td>
    <td><b>34.67 <span style="color:red;">†1.90</span></b></td>
  </tr>
</table>

Table 4: Performance of identified heads with a new set of general, math, and code datasets. The best results are highlighted in bold. † indicates change relative to the baseline. Our method consistently outperforms other methods, demonstrating the transferability of selected task heads. Full details are listed in Appendix G.

<table>
  <tr>
    <th>Model</th>
    <th>Method</th>
    <th>MMLU</th>
    <th>ARC-C</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.2-1B</td>
    <td>vanilla</td>
    <td>32.2</td>
    <td>32.8</td>
    <td>32.5</td>
  </tr>
  <tr>
    <td>w/ full</td>
    <td>28.14</td>
    <td>26.95</td>
    <td>27.55</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>27.80</td>
    <td>23.18</td>
    <td>25.49</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>27.86</td>
    <td>26.27</td>
    <td>27.07</td>
  </tr>
  <tr>
    <td><b>ALPS</b></td>
    <td><b>29.88</b></td>
    <td><b>28.73</b></td>
    <td><b>29.31</b></td>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.2-3B</td>
    <td>vanilla</td>
    <td>58</td>
    <td>69.1</td>
    <td>63.55</td>
  </tr>
  <tr>
    <td>w/ full</td>
    <td>44.22</td>
    <td>50.82</td>
    <td>47.52</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>46.16</td>
    <td>46.95</td>
    <td>46.56</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>45.55</td>
    <td>47.30</td>
    <td>46.43</td>
  </tr>
  <tr>
    <td><b>ALPS</b></td>
    <td><b>48.83</b></td>
    <td><b>52.37</b></td>
    <td><b>50.60</b></td>
  </tr>
  <tr>
    <td rowspan="4">Llama-3.1-8B</td>
    <td>vanilla</td>
    <td>66.7</td>
    <td>79.7</td>
    <td>73.2</td>
  </tr>
  <tr>
    <td>w/ full</td>
    <td>55.40</td>
    <td>60.34</td>
    <td>57.87</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>22.95</td>
    <td>49.01</td>
    <td>35.98</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>54.59</td>
    <td>62.92</td>
    <td>58.76</td>
  </tr>
  <tr>
    <td><b>ALPS</b></td>
    <td><b>57.87</b></td>
    <td><b>64.29</b></td>
    <td><b>61.08</b></td>
  </tr>
</table>

Table 5: Performance comparison on two common general benchmarks between our method and baselines with the vanilla models. Full details are listed in Appendix G.

ing specialized attention requirements for mathematical reasoning. Some heads appear consistently important across all three tasks, indicating partial head sharing that aligns with similarities in data distributions as illustrated in Appendix A. These findings highlight that our method effectively identifies both task-specific and task-agnostic heads, shedding light on how different downstream objectives influence attention-head specialization.

![Bar chart showing average performance (%) for ALPS and baselines on Llama-3.2-1B, Llama-3.2-3B, and Llama-3.1-8B](page_1012_1342_388_246.png)

Figure 6: Evaluation on transferability of ALPS with a new set of general, math, and code datasets. Results with † correspond to results from Table 1, while ‡ represents performance from Table 4.

4.5 Does Heads Transfer?

To assess whether ALPS-selected heads generalize across different datasets within the same task domain, we evaluate them on a new set of general, math, and code datasets, as illustrated in Section 4.1. Table 4 and Figure 6 show that our method consistently outperforms all baselines, indicating the transferability of the task-sensitive attention heads identified by ALPS that can be reused effectively on new datasets within the same tasks. This highlights the potential of our method for reducing alignment costs without sacrificing performance.

4.6 Fewer Heads Avoid Forgetting

Table 5 presents results on MMLU and ARC-C, two commonly used benchmarks to assess general knowledge, making them vulnerable to knowledge forgetting when the alignment dataset is limited in quality or quantity (Chang et al., 2024). The w/full baseline shows a noticeable drop in performance, reflecting the loss of pre-trained knowledge. In contrast, ALPS , which restricts updates to a small fraction of attention heads, better preserves this general knowledge and consistently outperforms both w/ full, w/o attn, and Random baselines. By limiting the number of trained heads that introduce sparsity, our method reduces overfitting to the smaller task-specific dataset, thereby mitigating knowledge forgetting from the pre-training phase.

5 Conclusion

This work presents ALPS , an approach that leverages model weight parameters to localize and prune task-sensitive attention heads, leading to more efficient LLM alignment while avoiding data dependency. Extensive experiments on general, math, and code tasks demonstrate that our method consistently outperforms baselines, exhibits head transferability, and mitigates knowledge forgetting. Overall, ALPS offers a promising perspective on efficient LLM alignment and paves the way for further research into parameter-efficient alignment strategies that capitalize on the role of attention heads.
6 Limitation

While ALPS demonstrates promising improvements in efficiency and task performance, our investigation into its transferability and its ability to mitigate knowledge forgetting remains preliminary. Due to training overhead constraints, our evaluations have been limited to a select set of downstream tasks and datasets. Moreover, further investigation is needed to elucidate the underlying mechanism by which ALPS preserves pre-trained knowledge. Despite these limitations, our approach still effectively avoids data dependency and highlights a promising direction for future research in efficient LLM alignment.

7 Acknowledgment

This work is mainly supported by the National Key Research and Development Program of China (No. 2022YFB3304100), and partially by the Fundamental Research Funds for the Central Universities (226-2024-00049).

References

Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. 2023. Gpt-4 technical report. arXiv preprint arXiv:2303.08774.

Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. 2023. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245.

Martin Arjovsky, Soumith Chintala, and Léon Bottou. 2017. Wasserstein gan. Preprint, arXiv:1701.07875.

Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. 2021. Program synthesis with large language models. arXiv preprint arXiv:2108.07732.

Yupeng Chang, Xu Wang, Jindong Wang, Yuan Wu, Linyi Yang, Kaijie Zhu, Hao Chen, Xiaoyuan Yi, Cunxiang Wang, Yidong Wang, et al. 2024. A survey on evaluation of large language models. ACM Transactions on Intelligent Systems and Technology, 15(3):1–45.

Sahil Chaudhary. 2023. Code alpaca: An instruction-following llama model for code generation. https://github.com/sahil280114/codealpaca.

Hao Chen, Yiming Zhang, Qi Zhang, Hantao Yang, Xi-aomeng Hu, Xuetao Ma, Yifan Yanggong, and Junbo Zhao. 2023a. Maybe only 0.5% data is needed: A preliminary exploration of low training data instruction tuning. arXiv preprint arXiv:2305.09246.

Jianhui Chen, Xiaozhi Wang, Zijun Yao, Yushi Bai, Lei Hou, and Juanzi Li. 2024. Finding safety neurons in large language models. arXiv preprint arXiv:2406.14144.

Lichang Chen, Shiyang Li, Jun Yan, Hai Wang, Kalpa Gunaratna, Vikas Yadav, Zheng Tang, Vijay Srinivasan, Tianyi Zhou, Heng Huang, et al. 2023b. Al-pagasaus: Training a better alpaca with fewer data. arXiv preprint arXiv:2307.08701.

Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, et al. 2021. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374.

Kevin Clark. 2019. What does bert look at? an analysis of bert’s attention. arXiv preprint arXiv:1906.04341.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. 2018. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457.

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. 2021. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168.

Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. 2023. Enhancing chat language models by scaling high-quality instructional conversations. Preprint, arXiv:2305.14233.

Yu Fu, Zefan Cai, Abedelkadir Asi, Wayne Xiong, Yue Dong, and Wen Xiao. 2024. Not all heads matter: A head-level kv cache compression method with integrated retrieval and reasoning. arXiv preprint arXiv:2410.19258.

Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Muennighoff, et al. 2024. A framework for few-shot language model evaluation.

Atticus Geiger, Hanson Lu, Thomas Icard, and Christopher Potts. 2021. Causal abstractions of neural networks. Advances in Neural Information Processing Systems, 34:9574–9586.

Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. 2025. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948.
Daya Guo, Qihao Zhu, Dejian Yang, Zhenda Xie, Kai Dong, Wentao Zhang, Guanting Chen, Xiao Bi, Yu Wu, YK Li, et al. 2024. Deepseek-coder: When the large language model meets programming—the rise of code intelligence. arXiv preprint arXiv:2401.14196.

Wes Gurnee, Neel Nanda, Matthew Pauly, Katherine Harvey, Dmitrii Troitskii, and Dimitris Bertsimas. 2023. Finding neurons in a haystack: Case studies with sparse probing. arXiv preprint arXiv:2305.01610.

Shwai He, Guoheng Sun, Zheyu Shen, and Ang Li. 2024. What matters in transformers? not all attention is needed. arXiv preprint arXiv:2406.15786.

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300.

Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874.

Geoffrey Hinton. 2015. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531.

Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2021. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685.

Binyuan Hui, Jian Yang, Zeyu Cui, Jiaxi Yang, Dayiheng Liu, Lei Zhang, Tianyu Liu, Jiajun Zhang, Bowen Yu, Kenning Lu, et al. 2024. Qwen2.5-coder technical report. arXiv preprint arXiv:2409.12186.

Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, et al. 2018. Quantization and training of neural networks for efficient integer-arithmetic-only inference. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2704–2713.

Goro Kobayashi, Tatsuki Kuribayashi, Sho Yokoi, and Kentaro Inui. 2020. Attention is not only a weight: Analyzing transformers with vector norms. arXiv preprint arXiv:2004.10102.

Solomon Kullback and Richard A Leibler. 1951. On information and sufficiency. The annals of mathematical statistics, 22(1):79–86.

Guohao Li, Hasan Abed Al Kader Hammoud, Hani Itani, Dmitrii Khizbullin, and Bernard Ghanem. 2023a. Camel: Communicative agents for "mind" exploration of large scale language model society. Preprint, arXiv:2303.17760.

Xian Li, Ping Yu, Chunting Zhou, Timo Schick, Omer Levy, Luke Zettlemoyer, Jason Weston, and Mike Lewis. 2023b. Self-alignment with instruction back-translation. arXiv preprint arXiv:2308.06259.

David Lindner, János Kramár, Sebastian Farquhar, Matthew Rahtz, Thomas McGrath, and Vladimir Mikulik. 2023. Tracr: Compiled transformers as a laboratory for interpretability. Preprint, arXiv:2301.05062.

Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Deng, Chong Ruan, et al. 2024. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434.

Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. 2023. Is your code generated by chat-GPT really correct? rigorous evaluation of large language models for code generation. In Thirty-seventh Conference on Neural Information Processing Systems.

I Loschilov. 2017. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101.

meta llama. 2023. Llama recipes: Examples to get started using the llama models from meta. https://github.com/meta-llama/llama-cookbook/tree/archive-main.

Paul Michel, Omer Levy, and Graham Neubig. 2019. Are sixteen heads really better than one? Advances in neural information processing systems, 32.

OpenAI. 2025. Openai o3-mini.

David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R Bowman. 2023. Gpqa: A graduate-level google-proof q&a benchmark. arXiv preprint arXiv:2311.12022.

Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, et al. 2023. Code llama: Open foundation models for code. arXiv preprint arXiv:2308.12950.

Robert Schmirler, Michael Heinzinger, and Burkhard Rost. 2024. Fine-tuning protein language models boosts predictions across diverse tasks. Nature Communications, 15(1):7407.

Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al. 2024. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300.

Guangyuan Shi, Zexin Lu, Xiaoyu Dong, Wenlong Zhang, Xuanyu Zhang, Yujie Feng, and Xiao-Ming Wu. 2024. Understanding layer significance in llm alignment. arXiv preprint arXiv:2410.17875.
Karan Singhal, Shekoofeh Azizi, Tao Tu, S Sara Mahdavi, Jason Wei, Hyung Won Chung, Nathan Scales, Ajay Tanwani, Heather Cole-Lewis, Stephen Pfohl, et al. 2023. Large language models encode clinical knowledge. Nature, 620(7972):172–180.

Karan Singhal, Tao Tu, Juraj Gottweis, Rory Sayres, Ellery Wulczyn, Mohamed Amin, Le Hou, Kevin Clark, Stephen R Pfohl, Heather Cole-Lewis, et al. 2025. Toward expert-level medical question answering with large language models. Nature Medicine, pages 1–8.

Derek Tam, Mohit Bansal, and Colin Raffel. 2024. Merging by matching models in task parameter subspaces. Transactions on Machine Learning Research.

Hanlin Tang, Yang Lin, Jing Lin, Qingsen Han, Shikuan Hong, Yiwu Yao, and Gongyi Wang. 2024. Razorattention: Efficient kv cache compression through retrieval heads. arXiv preprint arXiv:2407.15891.

Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. 2023. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/stanford_alpaca.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajwal Bhargava, Shruti Bhosale, et al. 2023. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288.

Leonid Nisonovich Vaserstein. 1969. Markov processes over denumerable products of spaces, describing large systems of automata. Problemy Peredachi Informatsii, 5(3):64–72.

Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. 2019. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:1905.09418.

Zhongwei Wan, Xin Wang, Che Liu, Samiul Alam, Yu Zheng, Jiachen Liu, Zhongnan Qu, Shen Yan, Yi Zhu, Quanlu Zhang, et al. 2023. Efficient large language models: A survey. arXiv preprint arXiv:2312.03863.

Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A Smith, Daniel Khashabi, and Hannaneh Hajishirzi. 2022. Self-instruct: Aligning language models with self-generated instructions. arXiv preprint arXiv:2212.10560.

Yuxiang Wei, Zhe Wang, Jiawei Liu, Yifeng Ding, and Lingming Zhang. 2023. Magicoder: Source code is all you need. arXiv preprint arXiv:2312.02120.

Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. 2024. Retrieval head mechanistically explains long-context factuality. arXiv preprint arXiv:2404.15574.

Can Xu, Qingfeng Sun, Kai Zheng, Xiubo Geng, Pu Zhao, Jiazhan Feng, Chongyang Tao, and Daxin Jiang. 2023. Wizardlm: Empowering large language models to follow complex instructions. arXiv preprint arXiv:2304.12244.

An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, et al. 2025. Qwen2.5 technical report. Preprint, arXiv:2412.15115.

An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jianhong Tu, Jingren Zhou, Junyang Lin, et al. 2024. Qwen2.5-math technical report: Toward mathematical expert model via self-improvement. arXiv preprint arXiv:2409.12122.

Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wen-hao Huang, Huan Sun, Yu Su, and Wenhua Chen. 2023. Mammoth: Building math generalist models through hybrid instruction tuning. arXiv preprint arXiv:2309.05653.

Haiyan Zhao, Hanjie Chen, Fan Yang, Ninghao Liu, Huiqi Deng, Hengyi Cai, Shuaiqiang Wang, Dawei Yin, and Mengnan Du. 2024. Explainability for large language models: A survey. ACM Transactions on Intelligent Systems and Technology, 15(2):1–38.

Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, et al. 2023. A survey of large language models. arXiv preprint arXiv:2303.18223.

Zifan Zheng, Yezhaohui Wang, Yuxin Huang, Shichao Song, Mingchuan Yang, Bo Tang, Feiyu Xiong, and Zhiyu Li. 2024. Attention heads of large language models: A survey. arXiv preprint arXiv:2409.03752.

Weishun Zhong, Ben Sorscher, Daniel Lee, and Haim Sompolinsky. 2022. A theory of weight distribution-constrained learning. Advances in Neural Information Processing Systems, 35:14113–14127.

Chunting Zhou, Pengfei Liu, Puxin Xu, Srinivasan Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, et al. 2024a. Lima: Less is more for alignment. Advances in Neural Information Processing Systems, 36.

Jeffrey Zhou, Tianjian Lu, Swaroop Mishra, Siddhartha Brahma, Sujoy Basu, Yu Lan, Denny Zhou, and Le Hou. 2023. Instruction-following evaluation for large language models. arXiv preprint arXiv:2311.07911.

Zhenhong Zhou, Haiyang Yu, Xinghua Zhang, Rongwu Xu, Fei Huang, Kun Wang, Yang Liu, Junfeng Fang, and Yongbin Li. 2024b. On the role of attention heads in large language model safety. arXiv preprint arXiv:2410.13708.

Andy Zou, Long Phan, Sarah Chen, James Campbell, Phillip Guo, Richard Ren, Alexander Pan, et al. 2023. Representation engineering: A top-down approach to ai transparency. arXiv preprint arXiv:2310.01405.
A  Datasets and Benchmarks

Datasets. Table 6 provides an overview of the datasets used in our experiments. Each dataset consists of \( \langle \text{instruction}, \text{answer} \rangle \) pairs, which serve as the basis for supervised fine-tuning. These datasets span diverse domains, ensuring a comprehensive evaluation of our approach across general, math, and code tasks

<table>
  <tr>
    <th>Name</th>
    <th>Task</th>
    <th># Samples</th>
  </tr>
  <tr>
    <td>UltraChat (Ding et al., 2023)</td>
    <td>general</td>
    <td>200k</td>
  </tr>
  <tr>
    <td>MathInstruct (Yue et al., 2023)</td>
    <td>math</td>
    <td>262k</td>
  </tr>
  <tr>
    <td>Magicoder (Wei et al., 2023)</td>
    <td>code</td>
    <td>110k</td>
  </tr>
  <tr>
    <td>Alpaca (Taori et al., 2023)</td>
    <td>general</td>
    <td>52k</td>
  </tr>
  <tr>
    <td>Camel-math (Li et al., 2023a)</td>
    <td>math</td>
    <td>50k</td>
  </tr>
  <tr>
    <td>CodeAlpaca (Chaudhary, 2023)</td>
    <td>code</td>
    <td>20k</td>
  </tr>
</table>

Table 6: Overview of the datasets employed in our experiments, including dataset name, task domain, and number of samples.

Data Distribution. Figure 7 shows a 3D visualization of the data distribution for UltraChat, MathInstruct, and Magicoder, as embedded by the Llama-3.1-8B model. Notably, the code and math datasets exhibit some overlaps, whereas the general dataset displays a distinct distribution, reflecting domain-specific characteristics.

![3D visualization of data distribution for UltraChat, MathInstruct, and Magicoder. All data is embedded using the Llama-3.1-8B model.](page_491_1047_484_384.png)

Figure 7: 3D visualization of data distribution for UltraChat, MathInstruct, and Magicoder. All data is embedded using the Llama-3.1-8B model.

Evaluations. We evaluate model performance using lm-eval (Gao et al., 2024; meta llama, 2023) to assess general capabilities such as instruction following, reasoning, dialogue, and mathematical abilities, and EvalPlus (Liu et al., 2023) to evaluate the capabilities of code generation. All settings are with the official report, and the details regarding the evaluation benchmarks, including the number of shots and metrics, are summarized in Table 7, with further metric-specific details provided in Table 8

B  Training Details

For all experiments, we use the Alpaca (Taori et al., 2023) template for fine-tuning across all datasets and models. Table 9 shows the training template employed, which standardizes the input format and instruction style for all tasks. This consistency facilitates effective fine-tuning across diverse domains.

C  Does ALPS Require a Full Fine-tuning Cycle to Obtain a Task-related Model?

For better clarification, the task-related model in our method serves merely as a recipe to extract the parameter shifts needed for head identification, and the primary focus of ALPS is on the identification of task-sensitive attention heads for subsequent parameter-efficient fine-tuning (PEFT), rather than on the process of obtaining a task-related model. For domains such as code and math, well-established task-specific models, e.g., CodeLlama (Roziere et al., 2023) and Qwen-2.5-Coder (Hui et al., 2024), are readily available and can be used directly in conjunction with the same size and architecture model, e.g., Llama-2 (Touvron et al., 2023), Qwen-2.5 (Yang et al., 2025) to perform ALPS head identification, and thereby obviating the need for additional full fine-tuning. Since there is no widely recognized task fine-tuned model with the Llama-3.2 series, we then opted to perform full SFT to obtain a task-related baseline. Importantly, this choice was made to ensure a fair comparison and does not affect the core contribution of ALPS, which lies in its efficient head identification strategy. Furthermore, transferable experiments in Table 4 demonstrate that once the task-sensitive heads are identified, they remain transferable across subsequent fine-tuning cycles. This one localization enables significant reductions in training cost for all downstream tasks, which serves as a key strength of ALPS.

To further support our claim that ALPS does not necessarily involve a full fine-tuning cycle to obtain a task-related model, we conducted experiments showing that even with 0-cost or low-cost training (10% parameter, 10% data size) to obtain the task-related model, ALPS still achieves com-
<table>
  <tr>
    <th>Benchmark</th>
    <th>Capability</th>
    <th># Shots</th>
    <th>Metric</th>
  </tr>
  <tr>
    <td>IFEval (<i>Zhou et al., 2023</i>)</td>
    <td>Instruction Following</td>
    <td>-</td>
    <td>Avg(Prompt/Instruction acc Loose/Strict)</td>
  </tr>
  <tr>
    <td>GPQA (<i>Rein et al., 2023</i>)</td>
    <td>Reasoning</td>
    <td>0</td>
    <td>acc</td>
  </tr>
  <tr>
    <td>GSM8K (<i>Cobbe et al., 2021</i>)</td>
    <td>Math</td>
    <td>8</td>
    <td>em_maj1@1</td>
  </tr>
  <tr>
    <td>MATH (<i>Hendrycks et al., 2021</i>)</td>
    <td>Math</td>
    <td>0</td>
    <td>final_em</td>
  </tr>
  <tr>
    <td>Humaneval (<i>Chen et al., 2021</i>)</td>
    <td>Code</td>
    <td>0</td>
    <td>pass@1</td>
  </tr>
  <tr>
    <td>Humaneval+ (<i>Liu et al., 2023</i>)</td>
    <td>Code</td>
    <td>0</td>
    <td>pass@1</td>
  </tr>
  <tr>
    <td>MBPP (<i>Austin et al., 2021</i>)</td>
    <td>Code</td>
    <td>0</td>
    <td>pass@1</td>
  </tr>
  <tr>
    <td>MBPP+ (<i>Liu et al., 2023</i>)</td>
    <td>Code</td>
    <td>0</td>
    <td>pass@1</td>
  </tr>
  <tr>
    <td>MMLU (<i>Hendrycks et al., 2020</i>)</td>
    <td>General</td>
    <td>5</td>
    <td>macro_avg/acc</td>
  </tr>
  <tr>
    <td>ARC-C (<i>Clark et al., 2018</i>)</td>
    <td>Reasoning</td>
    <td>0</td>
    <td>acc</td>
  </tr>
</table>

Table 7: Overview of benchmarks utilized in our experiments, detailing the assessed capability, number of shots, and evaluation metrics.

<table>
  <tr>
    <th>Metric</th>
    <th>Detail</th>
  </tr>
  <tr>
    <td>macro_avg/acc (Macro Average Accuracy)</td>
    <td>The mean accuracy across all classes or tasks</td>
  </tr>
  <tr>
    <td>acc_char (Character-Level Accuracy)</td>
    <td>Evaluates character-level correctness</td>
  </tr>
  <tr>
    <td>em (Exact Match)</td>
    <td>Measures how often the model output, exactly matches the reference answer</td>
  </tr>
  <tr>
    <td>f1 (F1 Score)</td>
    <td>Balances precision and recall, particularly useful for imbalanced data</td>
  </tr>
  <tr>
    <td>pass@1</td>
    <td>Assesses the correctness of generated code, on the first attempt</td>
  </tr>
  <tr>
    <td>em_maj1@1</td>
    <td>Measures exact-match on the first major attempt, especially in complex reasoning or math problems</td>
  </tr>
  <tr>
    <td>final_em</td>
    <td>The final exact-match score, commonly used in challenging benchmarks</td>
  </tr>
</table>

Table 8: Overview of evaluation metric details.

<table>
  <tr>
    <th>Field</th>
    <th>Content</th>
  </tr>
  <tr>
    <td>System prompt</td>
    <td>Below is an instruction that describes a task.<br>Write a response that appropriately completes the request.</td>
  </tr>
  <tr>
    <td>User prompt</td>
    <td>### Instruction: { {content} }<br>### Response:</td>
  </tr>
</table>

Table 9: Training template used for fine-tuning.

petitive performance, as shown in Table 10.

These results illustrate that training a model via head identification using an existing task-related model (0-cost) yields comparable performance, even though these readily available models may not have been fully trained. Moreover, a task model obtained with less data or fewer parameters (lower cost) also performs well after ALPS head identification, demonstrating that ALPS does not necessarily require a complete full fine-tuning cycle to obtain
<table>
  <tr>
    <th rowspan="2">Method</th>
    <th rowspan="2">\( \mathcal{M}_T \)</th>
    <th colspan="3">Math</th>
    <th colspan="3">Code</th>
  </tr>
  <tr>
    <th>GSM8K</th>
    <th>MATH</th>
    <th>HEval</th>
    <th>HEval+</th>
    <th>MBPP</th>
    <th>MBPP+</th>
  </tr>
  <tr>
    <th colspan="8"><i>Llama3.2-1B</i></th>
  </tr>
  <tr>
    <td>w/full</td>
    <td>SFT</td>
    <td>17.82</td>
    <td>4.60</td>
    <td>45.12</td>
    <td>39.08</td>
    <td>28.63</td>
    <td>23.78</td>
  </tr>
  <tr>
    <td>ALPS_Readily</td>
    <td>Readily \( \mathcal{M} \)</td>
    <td>21.81</td>
    <td>6.87</td>
    <td>45.78</td>
    <td>39.29</td>
    <td>27.93</td>
    <td>24.01</td>
  </tr>
  <tr>
    <td>ALPS_Fewer_\( \mathcal{D} \)</td>
    <td>10% data</td>
    <td>22.96</td>
    <td>7.06</td>
    <td>46.91</td>
    <td>39.87</td>
    <td>28.73</td>
    <td>23.62</td>
  </tr>
  <tr>
    <td>ALPS_Fewer_\( \theta \)</td>
    <td>attn only (10%)</td>
    <td><b>23.85</b></td>
    <td>7.19</td>
    <td><b>47.23</b></td>
    <td><b>41.01</b></td>
    <td>28.87</td>
    <td>23.98</td>
  </tr>
  <tr>
    <td><i>ALPS</i></td>
    <td>SFT</td>
    <td>23.74</td>
    <td><b>7.28</b></td>
    <td>46.89</td>
    <td>40.22</td>
    <td><b>29.13</b></td>
    <td><b>24.21</b></td>
  </tr>
  <tr>
    <th colspan="8"><i>Llama3.2-3B</i></th>
  </tr>
  <tr>
    <td>w/full</td>
    <td>SFT</td>
    <td>31.77</td>
    <td>16.68</td>
    <td>56.18</td>
    <td>50.53</td>
    <td>50.00</td>
    <td>41.50</td>
  </tr>
  <tr>
    <td>ALPS_Readily</td>
    <td>Readily \( \mathcal{M} \)</td>
    <td>33.28</td>
    <td>16.23</td>
    <td>58.27</td>
    <td>51.42</td>
    <td><b>54.87</b></td>
    <td><b>43.33</b></td>
  </tr>
  <tr>
    <td>ALPS_Fewer_\( \mathcal{D} \)</td>
    <td>10% data</td>
    <td>33.87</td>
    <td>16.98</td>
    <td>56.97</td>
    <td>49.68</td>
    <td>52.77</td>
    <td>41.98</td>
  </tr>
  <tr>
    <td>ALPS_Fewer_\( \theta \)</td>
    <td>attn only (10%)</td>
    <td>34.17</td>
    <td><b>17.33</b></td>
    <td>56.89</td>
    <td>50.05</td>
    <td>51.79</td>
    <td>40.32</td>
  </tr>
  <tr>
    <td><i>ALPS</i></td>
    <td>SFT</td>
    <td><b>34.27</b></td>
    <td>17.12</td>
    <td><b>58.28</b></td>
    <td><b>51.67</b></td>
    <td>52.21</td>
    <td>42.78</td>
  </tr>
</table>

Table 10: Experiments on variants of task-specific model acquisition.

the task-related model.

In summary, ALPS is a method focused on efficient task head identification. The task-related model serves merely as a component in the pipeline and the acquisition process does not necessarily require a full fine-tuning cycle, especially when leveraging readily available task models.

D Efficiency of ALPS

To better understand the efficiency of ALPS, we list the training time cost in Table 11. Our method shows significant reductions in training time, thus improving the efficiency.

<table>
  <tr>
    <th>Model</th>
    <th>Method</th>
    <th>Avg. Time Cost (hrs)</th>
  </tr>
  <tr>
    <td>Llama-3.2-1B</td>
    <td>w/ full</td>
    <td>2.02</td>
  </tr>
  <tr>
    <td></td>
    <td><i>ALPS</i></td>
    <td><b>0.40</b></td>
  </tr>
  <tr>
    <td>Llama-3.2-3B</td>
    <td>w/ full</td>
    <td>4.40</td>
  </tr>
  <tr>
    <td></td>
    <td><i>ALPS</i></td>
    <td><b>1.05</b></td>
  </tr>
  <tr>
    <td>Llama-3.1-8B</td>
    <td>w/ full</td>
    <td>7.67</td>
  </tr>
  <tr>
    <td></td>
    <td><i>ALPS</i></td>
    <td><b>2.18</b></td>
  </tr>
</table>

Table 11: Training time cost.

E Why is Parameter Alignment Distribution Score Better?

Let us assume that the change from \( W_{o,B}^h \) to \( W_{o,T}^h \) can be modeled as a small translation in the parameter space:

\[
W_{o,T}^h = W_{o,B}^h + \Delta,
\]

where \( \Delta \) is a small perturbation that reflects task-specific adjustments. The W1 distance between \( P_T^h \) and \( P_B^h \) is defined as:

\[
W_1 \left( P_B^h, P_T^h \right) = \inf_{\gamma \in \Gamma(P_B^h, P_T^h)} \mathbb{E}_{(x,y) \sim \gamma} \left[ \| x - y \| \right].
\]

For the special case where \( P_T^h \) is simply a translated version of \( P_B^h \), i.e., \( P_T^h \approx P_B^h(x - \Delta) \), leading to:

\[
W_1 \left( P_B^h, P_T^h \right) \approx \| \Delta \|.
\]

Thus, in the context of ALPS, \( s_{\mu}^{PAD} = W_1 (P_B^h, P_T^h) \) directly reflects the magnitude of the parameter shift caused by task-specific fine-tuning. This linear relationship ensures that even small shifts in the attention head's parameters are captured in proportion to \( \| \Delta \| \).

As for the KL divergence, the shifts between \( P_T^h \) and \( P_B^h \) is defined as:

\[
D_{KL}(P_B^h \| P_T^h) = \sum_i P_B^h(i) \log \frac{P_B^h(i)}{P_T^h(i)}.
\]

For small perturbations \( \Delta \), we can perform a second-order Taylor expansion. Under regularity conditions, this yields:

\[
D_{KL}(P_B^h \| P_T^h) \approx \frac{1}{2} \mathcal{I}(P_B^h) \| \Delta \|^2,
\]

where \( \mathcal{I}(P_B^h) \) is the Fisher information associated with \( P_B^h \).
Notice that while \( W_1 \) scales linearly with the magnitude of the shift \( \| \Delta \| \), the KL divergence scales quadratically. This quadratic behavior means that for small but significant shifts, KL divergence may underemphasize these changes compared to \( W_1 \) (Arjovsky et al., 2017). Moreover, KL divergence is asymmetric and can become unstable or even infinite if \( P^h_T \) assigns near-zero probability where \( P^h_S \) does not.

F  Heatmaps

Figure 8 presents heatmaps of the top 10% task-sensitive attention heads selected by our method for three Llama-3 models (Llama-3.2-1B, Llama-3.2-3B, and Llama-3.1-8B) across general, math, and code tasks. Each cell in the heatmap corresponds to a specific layer-head combination, and its color intensity reflects the \( s^{PAD} \) value. In the 1B and 3B models, the heatmaps for general and code tasks exhibit considerable overlap, whereas in the 8B model, the distributions for all three tasks are distinctly different. These observations indicate that as the model scale increases, the set of task-sensitive heads shifts, reflecting the diverse strategies that different models employ to handle task-specific requirements. Our method, ALPS, adaptively identifies these critical heads, demonstrating robust generalizability across various model scales and datasets.

G  Full Tables

In this section, we present the complete baseline results corresponding to those summarized in Table 4 and Table 5.

In Table 1 we mentioned that each task employs a separately fine-tuned model trained on its own dataset. With six baselines, three datasets, and three model scales, this setup requires \( 6 \times 3 \times 3 = 54 \) distinct models, making a full evaluation across ten benchmarks computationally prohibitive. Consequently, our main experiments report only per-task performance. To probe ALPS’s generalization further, we conducted selective evaluations of the baselines on additional tasks. Table 14 summarizes these results. These findings confirm that ALPS consistently generalizes better across diverse tasks.

<table>
  <tr>
    <th>Model</th>
    <th>Method</th>
    <th>General</th>
    <th>Math</th>
    <th>Code</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td rowspan="5">Llama-3.2-1B</td>
    <td>w/ full</td>
    <td>14.28</td>
    <td>5.87</td>
    <td>22.28</td>
    <td>14.14</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>19.73</td>
    <td>8.39</td>
    <td>19.98</td>
    <td>16.03 ± 1.89</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>21.80</td>
    <td>7.91</td>
    <td>19.63</td>
    <td>16.45 ± 12.31</td>
  </tr>
  <tr>
    <td>LC</td>
    <td>21.57</td>
    <td>8.02</td>
    <td>22.23</td>
    <td>17.27 ± 13.13</td>
  </tr>
  <tr>
    <td>LoRA</td>
    <td>22.98</td>
    <td>9.23</td>
    <td>23.27</td>
    <td>18.49 ± 4.35</td>
  </tr>
  <tr>
    <td></td>
    <td><b>ALPS</b></td>
    <td><b>23.43</b></td>
    <td><b>11.51</b></td>
    <td><b>24.11</b></td>
    <td><b>19.68 ± 5.54</b></td>
  </tr>
  <tr>
    <td rowspan="5">Llama-3.2-3B</td>
    <td>w/ full</td>
    <td>21.57</td>
    <td>15.23</td>
    <td>38.55</td>
    <td>25.12</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>22.31</td>
    <td>13.49</td>
    <td>35.72</td>
    <td>23.84 ± 1.28</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>22.38</td>
    <td>17.80</td>
    <td><b>41.37</b></td>
    <td>27.18 ± 12.06</td>
  </tr>
  <tr>
    <td>LC</td>
    <td>23.42</td>
    <td>17.20</td>
    <td>40.28</td>
    <td>26.97 ± 11.85</td>
  </tr>
  <tr>
    <td>LoRA</td>
    <td>24.01</td>
    <td>17.98</td>
    <td>40.57</td>
    <td>27.52 ± 12.40</td>
  </tr>
  <tr>
    <td></td>
    <td><b>ALPS</b></td>
    <td><b>24.65</b></td>
    <td><b>18.70</b></td>
    <td><b>40.24</b></td>
    <td><b>27.86 ± 2.74</b></td>
  </tr>
  <tr>
    <td rowspan="5">Llama-3.1-8B</td>
    <td>w/ full</td>
    <td>27.21</td>
    <td>28.76</td>
    <td>42.34</td>
    <td>32.77</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>18.29</td>
    <td>25.77</td>
    <td><b>46.11</b></td>
    <td>30.06 ± 13.38</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>27.93</td>
    <td><b>32.96</b></td>
    <td>40.43</td>
    <td>33.77 ± 11.00</td>
  </tr>
  <tr>
    <td>LC</td>
    <td>26.52</td>
    <td>30.28</td>
    <td>41.28</td>
    <td>32.69 ± 10.08</td>
  </tr>
  <tr>
    <td>LoRA</td>
    <td>27.87</td>
    <td>31.02</td>
    <td>42.27</td>
    <td>33.72 ± 10.95</td>
  </tr>
  <tr>
    <td></td>
    <td><b>ALPS</b></td>
    <td><b>28.21</b></td>
    <td>31.31</td>
    <td>44.57</td>
    <td><b>34.67 ± 11.90</b></td>
  </tr>
</table>

Table 12: Full table of Table 4.

<table>
  <tr>
    <th>Model</th>
    <th>Method</th>
    <th>MMLU</th>
    <th>ARC-C</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td rowspan="5">Llama-3.2-1B</td>
    <td>vanilla</td>
    <td>32.2</td>
    <td>32.8</td>
    <td>32.5</td>
  </tr>
  <tr>
    <td>w/ full</td>
    <td>28.14</td>
    <td>26.95</td>
    <td>27.55</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>27.80</td>
    <td>23.18</td>
    <td>25.49</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>27.86</td>
    <td>26.27</td>
    <td>27.07</td>
  </tr>
  <tr>
    <td>LC</td>
    <td>26.83</td>
    <td>27.08</td>
    <td>26.96</td>
  </tr>
  <tr>
    <td>LoRA</td>
    <td>28.79</td>
    <td>27.98</td>
    <td>28.39</td>
  </tr>
  <tr>
    <td></td>
    <td><b>ALPS</b></td>
    <td><b>29.88</b></td>
    <td><b>28.73</b></td>
    <td><b>29.31</b></td>
  </tr>
  <tr>
    <td rowspan="5">Llama-3.2-3B</td>
    <td>vanilla</td>
    <td>58</td>
    <td>69.1</td>
    <td>63.55</td>
  </tr>
  <tr>
    <td>w/ full</td>
    <td>44.22</td>
    <td>50.82</td>
    <td>47.52</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>46.16</td>
    <td>46.95</td>
    <td>46.56</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>45.55</td>
    <td>47.30</td>
    <td>46.43</td>
  </tr>
  <tr>
    <td>LC</td>
    <td>44.34</td>
    <td>48.98</td>
    <td>46.66</td>
  </tr>
  <tr>
    <td>LoRA</td>
    <td>47.28</td>
    <td>50.88</td>
    <td>49.08</td>
  </tr>
  <tr>
    <td></td>
    <td><b>ALPS</b></td>
    <td><b>48.83</b></td>
    <td><b>52.37</b></td>
    <td><b>50.60</b></td>
  </tr>
  <tr>
    <td rowspan="5">Llama-3.1-8B</td>
    <td>vanilla</td>
    <td>66.7</td>
    <td>79.7</td>
    <td>73.2</td>
  </tr>
  <tr>
    <td>w/ full</td>
    <td>55.40</td>
    <td>60.34</td>
    <td>57.87</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>22.95</td>
    <td>49.01</td>
    <td>35.98</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>54.59</td>
    <td>62.92</td>
    <td>58.76</td>
  </tr>
  <tr>
    <td>LC</td>
    <td>53.82</td>
    <td>61.89</td>
    <td>57.86</td>
  </tr>
  <tr>
    <td>LoRA</td>
    <td>56.28</td>
    <td>63.01</td>
    <td>59.65</td>
  </tr>
  <tr>
    <td></td>
    <td><b>ALPS</b></td>
    <td><b>57.87</b></td>
    <td><b>64.29</b></td>
    <td><b>61.08</b></td>
  </tr>
</table>

Table 13: Full table of Table 5.
Figure 8: Heatmaps of task-sensitive attention heads (10%) selected by our method on Llama-3.2-1B, Llama-3.2-3B, and Llama-3.1-8B, for **a**) general, **b**) math, and **c**) code tasks. The color intensity indicates the value of \( s_h^{PAD} \) of each head.
<table>
  <tr>
    <th rowspan="2">Model</th>
    <th rowspan="2">Method</th>
    <th rowspan="2">Task</th>
    <th rowspan="2">Attn %</th>
    <th colspan="3">General</th>
    <th colspan="2">Math</th>
    <th colspan="4">Code</th>
  </tr>
  <tr>
    <th>IFEval</th>
    <th>GPQA</th>
    <th>GSM8K</th>
    <th>MATH</th>
    <th>HEval</th>
    <th>HEval+</th>
    <th>MBPP</th>
    <th>MBPP+</th>
  </tr>
  <tr>
    <td rowspan="3">1B-MC</td>
    <td>w/ full</td>
    <td>code</td>
    <td>100%</td>
    <td>27.58</td>
    <td>10.27</td>
    <td><b>7.88</b></td>
    <td><b>2.50</b></td>
    <td><b>45.12</b></td>
    <td><b>39.08</b></td>
    <td><b>28.63</b></td>
    <td><b>23.78</b></td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>code</td>
    <td>0%</td>
    <td>28.30</td>
    <td>11.16</td>
    <td>7.28</td>
    <td>2.10</td>
    <td><b>41.51</b></td>
    <td><b>35.28</b></td>
    <td><b>29.91</b></td>
    <td><b>24.88</b></td>
  </tr>
  <tr>
    <td><i>ALPS</i></td>
    <td>code</td>
    <td>10%</td>
    <td><b>29.38</b></td>
    <td><b>12.27</b></td>
    <td>7.81</td>
    <td>2.48</td>
    <td><b>46.89</b></td>
    <td><b>40.22</b></td>
    <td><b>29.13</b></td>
    <td><b>24.21</b></td>
  </tr>
  <tr>
    <td rowspan="3">1B-MI</td>
    <td>w/ full</td>
    <td>math</td>
    <td>100%</td>
    <td>31.06</td>
    <td>8.26</td>
    <td><b>17.82</b></td>
    <td><b>4.60</b></td>
    <td>11.63</td>
    <td>10.42</td>
    <td>16.91</td>
    <td><b>16.17</b></td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>math</td>
    <td>0%</td>
    <td>30.18</td>
    <td>17.63</td>
    <td><b>23.35</b></td>
    <td><b>5.20</b></td>
    <td>12.87</td>
    <td>11.67</td>
    <td>16.92</td>
    <td>14.87</td>
  </tr>
  <tr>
    <td><i>ALPS</i></td>
    <td>math</td>
    <td>10%</td>
    <td><b>32.82</b></td>
    <td><b>18.72</b></td>
    <td><b>23.74</b></td>
    <td><b>7.28</b></td>
    <td><b>16.55</b></td>
    <td><b>14.07</b></td>
    <td><b>19.38</b></td>
    <td><b>15.68</b></td>
  </tr>
  <tr>
    <td rowspan="3">1B-UC</td>
    <td>w/ full</td>
    <td>general</td>
    <td>100%</td>
    <td><b>34.53</b></td>
    <td><b>18.75</b></td>
    <td><b>6.93</b></td>
    <td>1.24</td>
    <td>17.17</td>
    <td>14.62</td>
    <td>18.89</td>
    <td>14.83</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>general</td>
    <td>0%</td>
    <td><b>35.85</b></td>
    <td><b>19.64</b></td>
    <td>5.84</td>
    <td>2.52</td>
    <td><b>17.77</b></td>
    <td><b>15.21</b></td>
    <td>17.75</td>
    <td>12.72</td>
  </tr>
  <tr>
    <td><i>ALPS</i></td>
    <td>general</td>
    <td>10%</td>
    <td><b>36.69</b></td>
    <td><b>26.16</b></td>
    <td>6.44</td>
    <td><b>3.32</b></td>
    <td>16.55</td>
    <td>14.38</td>
    <td><b>19.32</b></td>
    <td><b>15.67</b></td>
  </tr>
  <tr>
    <td rowspan="3">3B-MC</td>
    <td>w/ full</td>
    <td>code</td>
    <td>100%</td>
    <td><b>40.41</b></td>
    <td><b>18.32</b></td>
    <td>28.73</td>
    <td>8.72</td>
    <td><b>56.18</b></td>
    <td><b>50.53</b></td>
    <td><b>50.00</b></td>
    <td><b>41.50</b></td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>code</td>
    <td>0%</td>
    <td>37.65</td>
    <td>17.41</td>
    <td>31.31</td>
    <td>9.44</td>
    <td><b>55.52</b></td>
    <td><b>50.62</b></td>
    <td>48.71</td>
    <td>39.72</td>
  </tr>
  <tr>
    <td><i>ALPS</i></td>
    <td>code</td>
    <td>10%</td>
    <td>37.65</td>
    <td>18.08</td>
    <td><b>32.26</b></td>
    <td><b>10.28</b></td>
    <td><b>58.28</b></td>
    <td><b>51.67</b></td>
    <td><b>52.21</b></td>
    <td><b>42.78</b></td>
  </tr>
  <tr>
    <td rowspan="3">3B-MI</td>
    <td>w/ full</td>
    <td>math</td>
    <td>100%</td>
    <td>35.25</td>
    <td>17.63</td>
    <td><b>31.77</b></td>
    <td><b>16.68</b></td>
    <td>26.28</td>
    <td>23.21</td>
    <td>29.92</td>
    <td>25.73</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>math</td>
    <td>0%</td>
    <td><b>35.37</b></td>
    <td>17.41</td>
    <td><b>33.21</b></td>
    <td><b>13.76</b></td>
    <td>28.27</td>
    <td>25.63</td>
    <td>31.72</td>
    <td>27.79</td>
  </tr>
  <tr>
    <td><i>ALPS</i></td>
    <td>math</td>
    <td>10%</td>
    <td>33.33</td>
    <td><b>18.33</b></td>
    <td><b>34.27</b></td>
    <td><b>17.12</b></td>
    <td><b>29.93</b></td>
    <td><b>26.82</b></td>
    <td><b>34.11</b></td>
    <td><b>28.82</b></td>
  </tr>
  <tr>
    <td rowspan="3">3B-UC</td>
    <td>w/ full</td>
    <td>general</td>
    <td>100%</td>
    <td><b>42.81</b></td>
    <td><b>22.32</b></td>
    <td>6.82</td>
    <td>8.16</td>
    <td>29.92</td>
    <td>27.43</td>
    <td>33.33</td>
    <td>29.13</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>general</td>
    <td>0%</td>
    <td><b>40.29</b></td>
    <td><b>24.33</b></td>
    <td>6.62</td>
    <td><b>8.22</b></td>
    <td><b>32.94</b></td>
    <td><b>28.49</b></td>
    <td><b>34.97</b></td>
    <td><b>28.82</b></td>
  </tr>
  <tr>
    <td><i>ALPS</i></td>
    <td>general</td>
    <td>10%</td>
    <td><b>44.96</b></td>
    <td><b>24.33</b></td>
    <td>5.98</td>
    <td>7.92</td>
    <td>29.33</td>
    <td>26.82</td>
    <td>34.43</td>
    <td><b>29.97</b></td>
  </tr>
  <tr>
    <td rowspan="3">8B-MC</td>
    <td>w/ full</td>
    <td>code</td>
    <td>100%</td>
    <td>47.84</td>
    <td>18.75</td>
    <td>46.63</td>
    <td>13.40</td>
    <td><b>69.50</b></td>
    <td><b>64.28</b></td>
    <td><b>63.21</b></td>
    <td><b>52.38</b></td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>code</td>
    <td>0%</td>
    <td>47.48</td>
    <td>19.20</td>
    <td><b>49.13</b></td>
    <td>10.22</td>
    <td><b>70.72</b></td>
    <td><b>65.88</b></td>
    <td><b>63.27</b></td>
    <td><b>52.55</b></td>
  </tr>
  <tr>
    <td><i>ALPS</i></td>
    <td>code</td>
    <td>10%</td>
    <td><b>49.96</b></td>
    <td><b>19.64</b></td>
    <td>48.98</td>
    <td><b>15.48</b></td>
    <td><b>72.68</b></td>
    <td><b>65.03</b></td>
    <td><b>63.67</b></td>
    <td><b>52.88</b></td>
  </tr>
  <tr>
    <td rowspan="3">8B-MI</td>
    <td>w/ full</td>
    <td>math</td>
    <td>100%</td>
    <td>36.69</td>
    <td>20.76</td>
    <td><b>63.00</b></td>
    <td><b>20.52</b></td>
    <td>30.53</td>
    <td>28.78</td>
    <td>28.23</td>
    <td>23.80</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>math</td>
    <td>0%</td>
    <td>26.38</td>
    <td>22.78</td>
    <td><b>64.06</b></td>
    <td><b>15.48</b></td>
    <td>32.38</td>
    <td>31.28</td>
    <td>30.13</td>
    <td>27.62</td>
  </tr>
  <tr>
    <td><i>ALPS</i></td>
    <td>math</td>
    <td>10%</td>
    <td><b>38.13</b></td>
    <td><b>24.11</b></td>
    <td><b>64.13</b></td>
    <td><b>22.48</b></td>
    <td><b>35.43</b></td>
    <td><b>32.33</b></td>
    <td><b>34.18</b></td>
    <td><b>28.37</b></td>
  </tr>
  <tr>
    <td rowspan="3">8B-UC</td>
    <td>w/ full</td>
    <td>general</td>
    <td>100%</td>
    <td><b>51.20</b></td>
    <td><b>25.22</b></td>
    <td>54.06</td>
    <td>15.12</td>
    <td>32.33</td>
    <td>28.72</td>
    <td>41.54</td>
    <td>34.92</td>
  </tr>
  <tr>
    <td>w/o attn</td>
    <td>general</td>
    <td>0%</td>
    <td><b>40.41</b></td>
    <td><b>26.16</b></td>
    <td>53.28</td>
    <td>14.28</td>
    <td>33.28</td>
    <td>28.62</td>
    <td>41.58</td>
    <td>34.78</td>
  </tr>
  <tr>
    <td><i>ALPS</i></td>
    <td>general</td>
    <td>10%</td>
    <td><b>51.33</b></td>
    <td><b>27.29</b></td>
    <td>54.28</td>
    <td><b>15.48</b></td>
    <td><b>35.42</b></td>
    <td><b>29.37</b></td>
    <td><b>42.33</b></td>
    <td><b>36.21</b></td>
  </tr>
</table>

Table 14: Selective baseline evaluation across additional tasks. Italicized and correspond to those in Table 1

Reviewer paper summaries: The paper introduces ALPS (Attention Localization and Pruning Strategy), a new method to make aligning large language models (LLMs) to downstream tasks more efficient. ALPS does this by finding and focusing on the most important attention heads for the task. Instead of updating all attention parameters, ALPS uses a parameter alignment distribution (PAD) score based on the Wasserstein-1 distance to locate these key heads and only updates them during training, which reduces computational costs. Experiments show that ALPS achieves performance improvements while updating just 10% of attention parameters, and it also demonstrates some degree of transferability across different datasets within the same task.

The paper presents ALPS (Attention Localization and Pruning Strategy), a novel method for improving the efficiency of Large Language Model (LLM) alignment by identifying and pruning task-sensitive attention heads. Instead of retraining all attention heads, ALPS selects the most critical heads using a parameter alignment distribution (PAD) score based on the Wasserstein-1 distance, thereby reducing computational overhead while maintaining or improving performance. The approach is evaluated across general, math, and code tasks on multiple Llama-3 model variants (1B, 3B, and 8B). Results demonstrate that ALPS updates only 10% of attention parameters while achieving up to 2.75% performance improvement over full fine-tuning baselines. Additionally, ALPS-selected heads exhibit transferability across datasets and help mitigate knowledge forgetting in LLM alignment.
Summary of strengths: - The paper introduces a PAD score-based method to identify critical attention heads, validated by extensive experiments showing improved performance with reduced computational costs.

- The study uses concise charts and figures to clearly demonstrate the effectiveness and impact of the proposed method.

1. The paper proposes an attention head pruning strategy that achieves competitive or superior performance while reducing alignment costs.
The parameter alignment distribution (PAD) score effectively quantifies task-induced shifts in attention heads, offering a principled way to identify task-relevant heads.

2. The ablation studies confirm that PAD score is superior to alternative metrics (cosine similarity, Euclidean distance, KL divergence) for attention head selection.

3. Extensive experiments on general, math, and code tasks across different Llama-3 model sizes (1B, 3B, 8B) provide robust validation of the method. ALPS outperforms full fine-tuning while using only 10% of attention heads, demonstrating its efficiency and effectiveness.
Summary of weaknesses: - This approach still relies on pre-fine-tuning a large model, potentially leading to substantial costs.

- What are the differences between the attention heads identified by this method and those obtained by activating model parameters with task-specific data, especially considering that the experiments were conducted separately for different task settings?

- The main experimental metrics do not significantly outperform LoRA, raising questions about the distinct advantages of ALPS.

- Since the $s^{PAD}$ metric does not show clear performance advantages over KL divergence, it still needs further exploration of other differences between them, such as their computational efficiencies.

- Tables 4 and 5 do not include comparisons with other baselines such as LoRA.

- The study does not consider out-of-domain scenarios, such as training on a general dataset and testing on math or code datasets, which is crucial for evaluating the model's generalization ability.

1. In introduction, the authors said that "However, these methods mainly rely on task-specific data to activate model parameters for attention head localization, inadvertently introducing data dependency to identified heads, which not only hinders generalization but also complicates reusability."  However, ALPS also needs a task fine-tuned model, which is trained on task-specific data from my understanding, to localize the important heads. Seems the overall computation cost of ALPS is even larger than related works. Please correct me if my understanding is wrong.
Comments, suggestions and typos: Refer to weaknesses.

There are some typos:

- The description from Line 109 to Line 112 is ambiguous and lacks clarity.

- In Line 176, it should read "n attention heads" instead.

1. line 019, maybe write it more specific.  "that our method activates only 10% of attention parameters **during fine-tuning**"
Overall assessment scores: 3

3
Metareview score: 3.5
Metareview: All reviewers found merit in this work while raising questions regarding the claimed advantages and some suggesting comparisons to additional baselines from the knowledge editing literature. I agree with the author response that such comparisons are not necessary for this contribution. The authors also clarified the reliance on a fully finetuned model and promised to update their manuscript to make the claims more accurate and to include additional helpful results.
Therefore I think the paper can be included in a *CL conference with these revisions.
Summary of reasons to publish: * The paper contributes a new method for parameter-efficient fine-tuning, which first determines which attention heads are most important for a specific task by using a fully-finetuned model for a related task (which could be fine-tuned on only 1% of the data), and then fine-tuning a model with all but 10% of the attention parameters frozen. 
* There are accuracy improvements over LORA and sometimes full fine-tuning. Therefore, the paper offers accuracy and efficiency gains compared to the baselines.