Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs
Authors: Qingru Zhang
Chandan Singh
Liyuan Liu
Xiaodong Liu
Bin Yu
Jianfeng Gao
Tuo Zhao
Venue: ICLR 2024 poster (2024)
Keywords: Attention steering
Post-hoc
Contextual emphasizing
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We propose PASTA, a post-hoc attention steering approach to enable users to hightlight specific information for LLMs and steers models to interpret the user-specified texts like human readers.

Abstract: In human-written articles, we often leverage the subtleties of text style, such as bold and italics, to guide the attention of readers. These textual emphases are vital for the readers to grasp the conveyed information.  When interacting with large language models (LLMs), we have a similar need -- steering the model to pay closer attention to user-specified information, e.g., an instruction. Existing methods, however, are constrained to process plain text and do not support such a mechanism. This motivates us to introduce PASTA -- Post-hoc Attention STeering Approach, a method that allows LLMs to read text with user-specified emphasis marks. To this end, PASTA identifies a small subset of attention heads and applies precise attention reweighting on them, directing the model attention to user-specified parts. Like prompting, PASTA is applied at inference time and does not require changing any model parameters. Experiments demonstrate that PASTA can substantially enhance an LLM's ability to follow user instructions or integrate new knowledge from user inputs, leading to a significant performance improvement on a variety of tasks, e.g., an average accuracy improvement of 22\% for LLAMA-7B. Our code is publicly available at https://github.com/QingruZhang/PASTA .

Published as a conference paper at ICLR 2024
TELL YOUR MODEL WHERE TO ATTEND: POST-HOC ATTENTION STEERING FOR LLMs
Qingru Z Chandan Singh', Liyuan Liu", Xiaodong Liu", Bin Yu', Jianfeng Gaod, Tuo Zhao
*Georgia Institute of Technology Universitty of California, Berkeley Microsoft Research {qingru. zhang, toutzhace@gaaecch.
binyubberkeley.. edu (chansingh, lucliu, xiaodl, ffqao cmicoooo
ABSTRACT
In human-written articles, we often leverage the subtleties of text style, such bold and italics, to guide the attention of readers. These textual emphases are vital for the readers to grasp the conveyed information. When interacting with large language models (LLMs), we have similar need steering the model to pay closer attention to user-specified information, e.g., an instruction. Existing methods, however, are constrained to process plain text and do not support such a mechanism. This motivates us to introduce PASTA Post-hoc Attention STeering Approach, method that allows LLMs to read text with user-specified emphasis marks. To this end, PASTA identifies a small subset of attention heads and applies precise attention reweighting them, directing the model attention to user-specified parts. Like prompting. PASTA applied at inference time and does not require changing any model parameters. Experiments demonstrate that PASTA can substantially enhance an LLM's ability to follow user instructions or integrate new knowledge from user inputs, leading to significant performance improvement on a variety of tasks, e.g., an average accuracy improvement of 22% for LLAMA-7B. Our code is publicly available at https: : g thub. com/ httgg/gggngggggggntara
such
INTRODUCTION
The advent of large language models (LLMs) has marked significant milestone in natural language processing (NLP) and artificial intelligence (Al), showcasing exceptional performance across a wide range of tasks (Vaswani al., 2017; Brown al., 2020a; OpenAI, 2023). Efforts to further refine these models have been relentless. aiming to enable them to process and respond to natural and programming languages with human- expertise (Stiennon et al., 2020; Yao et al., 2023).
Despite their remarkable achievements, LLMs often encounter challenges understanding their contextual inputs during interactions with users (Shen et al., 2023; Lu et al., 2021). This difficulty becomes particular evident when they are presented prompts containing extensive background contexts complex user instructions. Lengthy contexts can overwhelm LLMs, as their attention modules, learned from data, are unable to fully capture crucial details (Liu et al., 2023). Complex instructions can further inhibit the model from focusing the user's 's intentions, resulting in undesired outputs (Wei et al., 2022). Additionally, for time-sensitive data, such as news articles, there can exist factual knowledge within contexts, which contradicts with model prior beliefs induced from outdated pre-training. As a result, a model may generate outputs conditioned on its pre-existing belief instead of attending to new facts within the contexts (Meng al., 2022a;b; Mitchell et al., 2022: Hernandez et al., 2023). All of these challenges contribute to LLMs struggling to comprehend user intentions.
Compared to LLMs, human readers rarely struggle to understand the emphases of articles and intentions of writers. Writers often leverage variety of text styles, such bold and italies, to emphasize specific contents. This mechanism enables writers to direct and maintain the attention of
Work is done during Qingru Zhang's internship at Microsoft Research. We use prompts to refer to all LIM text inputs, including user instructions, and the other background information (which refer to context).

---

Published as a conference paper at ICLR 2024
Original user input:
Steered user input: Mary is doctor... "Return her occupation in json format"
Pay attention to Jon format"
Mory is doctor... Return her occupation in |son format
Oo
Layer 1
For each selected head:
Layer
LLM
Emphasize attention score for selected token positions
Layer
Head 1
Head
Head
Original output: X
Steered output:
Mory is working professional
("Name": "Mary", "Occupation": "Doctor")
Figure PASTA uses a user-specified part of the input to steer the model generation aligning with user intentions. PASTA modifies the attention scores generated during inference, by emphasizing the scores generated token positions corresponding to the user-specified part of the context.
human readers, ensuring that the intended information is accurately captured. In interactions between users and LLMs, it is users also need to highlight specific information for the model. Consequently, model generation can be effectively biased in accordance with user guidance, thus addressing the challenges mentioned earlier. This feature is particularly essential when designing user-AI interfaces, and can be frequently applied in extensive conversations between users and models. Existing methods, however, do not support such mechanism. L.MMs inherently limited processing plain texts, devoid of any stylistic cues or emphasis markers (Brown et al., 2020b; Liu et al., 2021; Wei et al., 2022). Even when emphasis markers are added to prompts, state-of-the-as LLMs often struggle to discern weak signals from couple of marker tokens (See evidence in Section 5.1).
Motivated by the need to convey user emphasis, we introduce PASTA (Post-hoc Attention STeering Approach), a post-hoc method that enables users highlight specific information, e-g., an instruction as in Figure/ 1, and steer models to interpret emphasized texts like human readers. Specifically, PASTA selects a small subset of attention heads and applies precise attention reweighting on them. As illustrated in Figure 1, PASTA upweights the attention scores of the user-specified tokens while downweighting the other tokens at specific attention heads. Our method is inspired by the observation that attention modules exhibit various token aoken-attndding patterns across different heads (Michel et al., 2019; Voita et al., 2019; Clark et al., 2019). These attention patterns can be interpreted as encoding diverse semantic
syntactic information, and altering them substantially influence model behaviors (Shi et al., 2023a; Hu et al., 2021b). Through steering attention modules, PASTA directs the model to pay close attention to the user-specified parts and hence generate the desired output aligning with the highlighted contents. Notably, PASTA is applied after training and does not require changing any model parameters; PASTA only requires access to the attention scores of specific heads of an LLM.
Since attention heads can serve different functions (Tenney et al., 2019; Deb et al., 2023), we introduce an efficient model profiling algorithm to identify which heads are effective for steering. Specifically, we subsample small training sets from multiple tasks and evaluate the performance of attention steering for each individual head across these tasks. PASTA selects the attention heads that, when steered, generally improve the multi-task performance. We empirically observe that steering these heads not only benefits the existing tasks but also enhances the performance on unseen tasks. Notably, the model profiling is performed only once for LLM. The selected attention heads can be regarded as a model-level profile, effective for steering the LLM on unseen tasks.
We conduct experiments on diverse tasks to demonstrate the effectiveness of PASTA. Specifically, evaluate PASTA using GPT-J-6B (Wang & Komatsuzaki, 2021) and LLAMA-7B (Touvron et al., 2023) on tasks that span complex instructions, lengthy contexts, and knowledge conflicts within contexts. The results demonstrate that PASTA consistently provides significant performance improvement over baseline prompting strategies. For example, PASTA achieve an average accuracy improvement of 22% over few-shot prompting for LLAMA-7B across challenging tasks.
Post-hoc means that method does update the model weights.
2

---

Published as a conference paper at ICLR 2024
2 BACKGROUND
Problem description In standard LLM prompting, given pre-trained LLM and text prompt z. In our setting, we additionally require (i) access to attention scores produced by attention modules in the LLM and (ii) we are provided user-specified subset of the prompt xg æ to be emphasized.
As in the example in Figure 1, æ can be a string that ends in an instruction, such as Mary is doctor but used to be netue...Return her occupation json format. If user emphasizes the instruction, 2g can simply be the final instruction Return her occupation in json format. In evaluation datasets, we assume that the user-specified part of each example is already provided by enclosing at its both ends in some emphasis markers, like "x" marker in Markdown. Generating these well-structured data often incurs little overhead. For example, in the dataset tailored for evaluting model ability to follow user instruction, can simply mark the final instruction for every example, which fixed and shared across examples. When it comes to user-LLM interface, users can specify ag by enclosing with same emphasis markers. Bg can be specified flexibly. Namely, it need not be continuous span, and can be used to emphasize diverse information.
Multi-Head Attention. A typical transformer model consists of stacked layers, where each layer contains two submodules: a multi-heat attention (MHA) and a fully connected (FFN). Given the input X E Raadd MHA of the layer performs the attention function in parallel H heads: MHA (1) (X) Concat(H ...., H( W. where
network
H V = Softmax (QK dh V
where Q XW qu, XWE XW., and W projection matrices of head h. d, is typically set to d H. Specifically, denote the attention scores at the head h of the I-th layer as A (1.b)
Wa: Ws, W., E RRdxd, are learnable
METHOD
PASTA (Algorithm 1) consists of two components: (i) post-hoc attention steering, which emphasizes the user-specified parts of the input during inference, Section 3.1 and (ii) multi-task model profiling, which selects the effective attention heads for steering, see Section 3.2.
Algorithm 1 PASTA: Post-hoc Attention Steering Approach
Multi-task model profiling (Section 3.2)
Input: small training sets (D(*) 1 1. the hyperparameters a, k;;
2: for < < m do
3: for 1 1 L, < do
4: Evaluate the model performance ((+) when steering the head (1, by (2);
5: Return the evaluation result of steering (1, h) on D(();
6: end for
7: Collect the steering results of all heads and return the task profiling R());
8: end for
9: Output: The attention head set 7( = 7 1-k"
Inference-time steering (Section 3.1)
1: Input: text inputs æ. user-underlined segments G, coefficient o;;
2: Output: the model generations while steering every head (1, /)) in H by (2).
3.1 POST-HOC ATTENTION STEERING
PASTA emphasizes the user-specified input subset by downweighting the attention scores of tokens that are not specified by the user. Specifically, given the index set highlighted input spans as G, PASTA emphasizes these user-specified tokens by an attention projection T:
aAgi/Ca j
H (1.b) (LA), where [T(A)ss =
(2)
A
otherwise.
We do need access model weights nor intermediate outputs from the other modules like PFNs.
3

---

Published as a conference paper at ICLR 2024
where 0 I a < 1 is scaling coefficient and G = [n] - G is the index set of tokens that are not in g. The term C, = Ejeg Aij + Ejeg- a,,,, normalizes the scores so that they sum to one. The attention steering (2) is conducted during the inference time and does not require any training.
(2) steers the model attention by scaling down the scores of tokens that are not highlighted by the user. When the coefficient is set very small, user-specified segments are highlighted given their increased attention scores after renormalization. Consequently, we can direct the model to concentrate more on the user-specified tokens, biasing the generation align with the specified contents.
PASTA scales down the attention scores of non-specified tokens by a. As renormalization is followed, is equivalent to scaling up the attention scores of user-specified tokens by 1/a. The reason of selecting (2) is that it can be more numerically stable compared to scaling up scores. Alternatively, one can also scale the attention scores by adding a positive constant to the underlined tokens G. The reason of we select multiplication in (2) instead of addition that it preserves the difference on attention magnitude among the highlighted tokens. As such, the steering operation only adjusts overall attention scales of two groups of tokens. In contrast, addition by large constant to the highlighted tokens results in their attention scores almost uniformly distributed, leading to unnecessary information loss and performance degeneration.
3.2 MULTI-TASK MODEL PROFILING
Empirically, we find that applying attention steering in (2) to all attention heads performs worse than applying it only to specific heads (see Section 5.3). It is important to specify the correct attention heads, given that different heads serve distinctive roles in encoding semantic/syntactic information. To this end, we propose multi-tast task model profiling algorithm to identify the effective attention heads for steering. Specifically, given tasks involving user emphases, we subsample small training set D(0) (e.g.. [DO 1000) from each task i. Then, we evaluate the performance of steering every individual attention head h) (1 S 1 s L, < h S H) on each small subset D (1 S i S m). For every task i, we rank all of heads according to their steering performance on D and regard the ranking R = [(42 ft)), (12, h2)................... as the profiling of task i. We then set the attention head set H. for steering as the intersection of top-k performing heads, H 1 for alternative choices). Intuitively, we expect performance to improve as the number of tasks m increases.
Ri(() (see Section 5.3
Like attention steering, model profiling requires only access to attention scores, in addition to its inputs and outputs (model weights and gradients are not required). Importantly, this process needs to be performed only once for I.I.M, similar to finetuning. However, unlike finetuning, model steering does not modify model weights and, more importantly, generalizes to new tasks. The resulting head set H. can be regarded as a model-level profile. Once it is determined, we can apply the attention steering on H to both existing tasks and unseen tasks to enhance model contextual understanding and benefit downstream performance.
4 EXPERIMENTAL SETUP
Evaluation tasks and metrics. We implement PASTA for two pre-trained models: GPT-J (6 billion parameters, (Wang & Komatsuzaki, 2021)) and LLaMA-7B (7 billion parameters, (Touvron et al.. 2023)). We evaluate the effectiveness of PASTA at (i) handling complex user instructions, (ii) interpreting lengthy contexts, and (iii) resolving in-context knowledge conflicts. For (i), we introduce two new tasks: JSON formatting and Pronouns changing. For (ii) and (iii), we study Bias in Bios (De-Arteaga et al., 2019) and CounterFact (Meng et al., 2022a). For each task, we provide a description, describing which part of the input we emphasize, and what metrics we use for evaluation (see Appendix A for full dataset details).
JSON Formatting is a new task that evaluates an LLM's ability to produce outputs in a user- desired format (JSON). This is an important usecase for LLMs when their output is being used in a downstream process. This task utilizes the biographical data from BiasBios (described below) but appends different instruction the end of contexts: answer the occupation of {person) and generate the answer as JSON format. The instruction prompts models to generate outputs in ISON format.
t We emphasize the final instruction
4

---

Published as a conference paper at ICLR 2024
Metrics: (a) Format aceuracy (F. Acc.) measures the accuracy generating valid JSON. (b) Prediction accuracy (P. Acc.) measures the accuracy at generating the correct target in JSON values after loading the JSON-formatted generations.
Pronouns changing is task that evaluates LLM's s ability to follow difficult user instruction. again uses the biographical contexts from BiasBios but instead instructs models to: substitute "she' and 'he" with 'they' and generate the occupation of (person) after changing pronouns.
& We emphasize the final instruction.
Metrics: (a) Accuracy evaluates the ratio that 'she/he' are successfully changed to 'they" in model generations. (b) All-changed accuracy (A. Acc.) is the ratio that models replace all corresponding pronouns, i.c., changing sherhheher/himmherrshiss to they/hhem/therr/hhiis....
CounterFact measures an LLM's ability to generate text consistent with new fact. Each example consists of (subject, relation, old target, new target), e.g., (Kevin Garnett, is a professional, basketball player, baseball player). We present the model both old and new facts following the prompt: Previously, {old fact), but currently, {new fact). {question). This change in facts over time often confuses LLMs, resulting random guesses on two of them when answering the {question).
/ We emphasize the input span containing the new fact.
Metrics: we evaluate metrics following (Meng et al.. 2022a): (a) Efficacy score (ES) is the portion of cases for which the model has PLLM target) > PLLM target); (b) Paraphrase score (PS) is the same as ES but changes the {question) with set of rephrased questions to assess the generalization
BiasBios consists of professional biographies of non-famous people, originally introduced to investigate gender bias occupations. Each example includes biographical context and a label of target occupation. The first sentence mentions the person's 's occupation, and subsequent sentences describe the individual's career history may not be directly related to the prediction, potentially distracting the model attention. At the end of the context, we append the question: {person} has the occupation of _.
/ We emphasize the first sentence, as it carries the most information about the occupation.
Metrics: following (Hernandez et al., 2023), we compute Accuracy by checking whether the probability assigned to the target occupation is the highest among the 28 candidate occupations.
For Pronouns changing, CounterFact, and BiasBios, additionally measure Fluency as the average bi-gram and tri-gram entropy of generations, designed to be low for degenerated or repetitive texts (Meng al., 2022a). We filter out any results receiving fluency below 3.0 (see full results including fluency in Appendix B.
Baselines.
We compare PASTA to the following baselines:
Zero-shot prompting is the most common approach to interact with LLMs, in which user feeds models a prompt containing background context and a user instruction or question.
Marked prompting alters the prompts used in zero-shot prompting by surrounding user-specified input spans with emphasis markers, e.g. asterisks, as is done markdown files for emphasis, or quotes, as is done in natural languages.
Few-shot prompting includes demonstrations (example inputs and target outputs) at the beginning of the prompt fed the LLM. Few-shot prompting often improves performance in new tasks, but increases the computational cost of inference due to the increased prompt length, particularly when demonstrations are lengthy (Dong et al., 2023); here we use 3 demonstrations context.
PASTA settings We study PASTA in 2 settings: multi-task and task-agnostic. In the multi-task setting, the evaluation task j is included for profiling, whereas the task-agnostic setting, the evaluation task is excluded (instead, we profile on the 3 datasets besides j). The multi-task setting improves performance but requires labeled training samples for the task which evaluated, which can be difficult to obtain in practice.
find that PASTA is not sensitive to the scaling coefficient (see Section 5.3) and fix it to 0.01 in our experiments. We select 1000 training samples from each the tasks above for model profiling. After model profiling, select l from (300, 400, 500) for LLAMA-7B
Empirically,
5

---

Published as a conference paper at ICLR 2024
Table 1: Main results LLAMA-7B to demonstrate that PASTA can improve model ability (i) follow user instruction (JSON Format and Prons. Changing); (ii) interpret contextual information (BiasBios); (iii) resolving knowledge conflicts (CounterFact). For all scores, higher is better. The best results are in bold.
JSON Format
Prons. Changing
BiasBios Acc
CounterFact
All
Method
F. Acc P. Acc
Acc / A.Acc
ES PS
Ave.
Zero-shot
60.00 / 54.94
71.84 66.28
87.36
58.50/ / 52.03
67.29
+markedd
18.55 / 12.71
39.14 35.17
90.62
57.74 / 50.52
49.38
Prompting
"m-mkeedd
4.56 4.20
20.55 18.19
89.82
58.14 51.70
42.15
Few-shot
84.85 / 73.58
59.06 55.27
88.79
87.45 / 49.82
73.45
Task-agnostic
88.16 / 49.08
83.65 81.31
93.54
98.82/ 99.03
85.89
PASTA
Multi-task
96.64 / 85.09
96.42 / 95.84
95.28
99.60 / 99.57
95.46
Table 2: Main results of GPT-J to demonstrate that PASTA can improve the model ability to (i) follow user instruction (JSON Format and Prons. Changing); (ii) interpret contextual information (BiasBios); (iii) resolving knowledge conflicts (CounterFact). For all scores, higher better. The best results are bold.
JSON Format
Prons. Changing
BiasBios Ace
CounterFact
All
Method
F. Ace P. Acc
Acc/ A.Ace
/ PS
Ave.
Zero-shot
28.83 / 25.09
39.88 36.19
72.76
42.14/ 42.02
44.96
marked
4.44 4.10
41.25 37.57
74.14
44.50/ 45.09
40.63
Prompting
"" ""-marked
8.81 5.62
6.12 5.72
78.64
45.54 41,84
33.87
Few-shot
84.15 / 72.65
35.77 32.08
72.98
68.34 / 38.23
59.65
Task-agnostic
46.68 / 34.71
91.62 88.60
80.84
99.54 99.57
77.80
PASTA
Multi-task
91.50 18.63
92.96 /9.334
94.96
98.62 98.79
85.22
to have the number of steered heads | {25, 53, 86). We find that PASTA achieves the best
performance on I.AMA-77 when < HH| < 100, i.e., k 400 or k = 500. For GPT-J, we select k from {250, 275, 300, 350} have [H| as {52, 72. 111, 153). For every task, split data into train/validationn/ess sets following (Hernandez et al., 2023) (See Appendix A) and select [H] by cross validation. For all tasks, model outputs generated with greedy search.
5
RESULTS
5.1
MAIN RESULT: PASTA IMPROVES MODEL GENERATION
Tables and present the main results for PASTA applied to LLAMA-7B and GPT-J respectively. Few-shot prompting is the strongest baseline, and task-agnostic PASTA outperforms it the main metric for each task for all settings except JSON Formatting with GPT-J. Multi-task PASTA outperforms all baselines across all settings.
PASTA improve LLM instruction following. The results from JSON Formatting and Pronouns Changing tasks indicate that, by highlighting the user instruction at the end of inputs, PASTA effectively steers models focus on user intentions. thereby biasing their generation to fulfill specific requirements or formats. For example, while GPT-J only achieves 39.9% of its zero-shot generations complying the user requirement on the Pronouns Changing task, PASTA yields remarkable 53% accuracy improvement by emphasizing the instruction. Moreover, PASTA achieves an impressive 96.64% format accuracy and 85.09% prediction accuracy when applied to LLAMA-7B on the JSON Formatting task. This performance exceeds that of few-shot prompting by 1 11% even though few-shot prompting explicitly provides the model with correct JSON examples through additional demonstrations. Table presents few examples generated by LLAMA-7B when applying PASTA.
PASTA can help models capture crucial contextual information. In the of BiasBios and CounterFact tasks, we apply PASTA to emphasize specific context spans for LLMs. Consequently, the models are guided to pay close attention the specific contextual information new facts within

---

Published as a conference paper at ICLR 2024
Table 3: Generation examples of I.AMM-77B on JSON Formatting and Pronouns Changing tasks.
Task
Prompt
Zero-shot Generation
PASTA Generation
JSON
Winnie is un American pholographer living New York. Specialized fash- ion photography and portrait, she ap- plies her talent both humans and mals instruction)
Winnie is a fashion and portrail pho- tographer living New York, special- ized in fashion photography and portrais, who applics her talent both homans and animals.
("name": "Winnie", "occupation". "pholographer"" }
Format
Prons.
Mary an Associate Professor Department of Curriculum Instruction at John University; she holds doc- torate in Bellling/Wiiigg/iiiooo the University of Peumsylvanin She re- cently published article "Urban Education" with... finttrcction)
Mary is a professue ment of Curiculum Instraction at St. John University; she holds dector- ste Reading/Wriiinn Writing/Litrrae from the University of Penmsylvaniu. She re- cently published article in "Urban Education" with...
Depart-
associate profes- sor in the department of curriculum and instruction at St. John's Univer- sity: they hold dectorate in read- ing/writin//lierraa from the University of Pennsylvania. They recently pub- lished an article in...
Mary
Change
from
contexts. The results from these two tasks illustrate that PASTA can direct the models to interpret the crucial information or resolve the knowledge conflicts within contexts, yielding significant improvement in prediction performance of both tasks. For example, PASTA achieves prediction accuracy of 94.96% for GPT-J on the BiasBios task, which is 16.32% higher than the best baseline.
Tables and 2 also suggest that marked prompting, a baseline that highlights specific texts akin to human writers, struggles to effectively convey emphasis to LLMs. One possible reason is that these emphasis markers rarely appear in the massive pre-training data. In contrast, few-shot prompt- ing sometimes leads to improvements in model performance. However, drawback of few-shot prompting is its instability, i.e. its performance exhibits high variance across different samples in the demonstration (See Appendix B).
5.2 PASTA CAN MITIGATE THE SENSITIVITY OF PROMPTS
Table 4: Results about sensitivity of model performance prompt rephrasing on the JSON Formatting task. Given rephrased instructions prompt template, PASTA can imporve zero-shot performance for all prompts.
LLAMA-7B
GPT-J
Instruction
Method
JSON Format F. Acc P. Ace
Prons. Changing Acc A. Acc
JSON Format
Prons, Changing Ace / Ace
Average
Acc Ace
Zero-shot
60.0 / 54.9
71,8 66.3
28.8 25.11
39.9 36.2
Original
PASTA
96.6 / 85.1
96.4 / 195.8
91.5 / 18.6
93.0
83.5
Zero-shot
36.0 32.4
49.2 42.6
25.4/ 17.1
56.5 / 54.8
39.3
Shortened
PASTA
87.4/ 65.9
89.0 86.9
54.1 /37.0
94.0/93.7
76.0
Zero-shot
57.9 / 54.2
82.3/ 79.6
76.0 72.8
67.1
Rephrased
PASTA
97.1/ 87.1
89.6 89.0
77.5 68.1
94.8 92.3
86.9
is well-known that the the performance of LLMs can be sensitive to minor changes in prompts, such rephrasing and reformatting, even when these prompts convey the same meaning (Reynolds & McDonell, 2021; Liu et al., 2021). We find that PASTA can alleviate the sensitivity of model performance varying prompts. Specifically, Table 4 evaluates the performance of LLAMA-7B and GPT-J on JSON Formatting and Pronouns Changing task given different instructions in the prompt template, all of which convey the same meaning (see precise prompts in Appendix A.1). The results show that zero-shot performance is sensitive to different prompts and can significantly deteriorate with poorly crafted templates. In contrast, PASTA consistently improves model performance over zero-shot prompting for all prompts, effectively mitigating sensitivity to variations in the prompts.
5.3 ANALYSIS AND ABLATIONS
In this section, investigate different hyperparameter choices and modeling decisions that affect the performance of PASTA.
Model profiling Figure 2 presents the results on the importance of model profiling introduced in Section 3.2. We compare PASTA when steering the selected heads versus other reasonable choices: steering (i) all heads, (ii) entire layers, or (iii) individual heads on the ISON Formatting task (See
7

---

Published as a conference paper at ICLR 2024
100
A.Acc
50
Zero-shot
25
Sleer all heads Steer entire layer
0
single
12
16 Layer
20
24
2B
32
Figure 2: The performance of LLAMA-7B on the ISON Formatting task when we steer (i) all heads (green); (ii) an entire layer (yellow); and (iii) an individual head within layer (blue violin plot). The performance varies dramatically across layers and across heads of a layer.
Appendix B.3 for comparisons on the remaining tasks). Selecting heads via model profiling in PASTA (red line) significantly outperforms other approaches. Steering all heads (dashed green line) degrades performance compared to the baseline zero-shot performance (dashed black line). This is likely because steering all heads over-amplifies the user-specified information the expense of other essential information required for effective generation and prediction. Interestingly, we find that the performance varies significantly when steering different layers (yellow) heads (blue violin plot). As mentioned in Section I, attention heads play distinct roles in encoding diverse semantic and syntactic information (Tenney et al., 2019). When steering heads, which are appropriately involved in encoding of user-specified information, the model can be guided to capture and reinforce these specific signals. Conversely, modifying the attention of unrelated heads not only fails to emphasize the desired information but also interferes with their original functions, resulting in performance deterioration. Therefore, it is important to identify the effective heads through model profiling prior to applying the steering.
Varying strategies for selecting heads during profiling. As described in Sec. 5.3, model profil- ing selects the Intersection of the top-k performing heads to steer across multiple tasks. Alternatively, when evaluating task can select heads for steering with different strategies: (i) Task Task-specific-
steer the top-k2 performing heads of only the task j. i.c., B or (ii) Union - the union these
[2(() heads across multiple tasks, i.c., - l:kg Table compares their performance. Using task-specific heads rather than crrreection-selcctee heads sometimes yields improved performance, but requires selecting different set heads for each new task.
Table 5:
Varying head selection strategies between top task-specific heads, union across multiple tasks, and intersection (the default used in PASTA),
JSON Format
Prons. Changing
BiasBios Acc
CounterFact ES / PS
All Avg.
PASTA
F. Acc / P. Acc
Acc A.Ace
Task-specific
95.56/ / 86.83
98.52 / 98.02
97.62
99.18 99.24
96.57
LLAMA
Union
88.42 / 74.49
92.12 / 91 ,44
96.36
99.24 99.35
92.22
Intersection
96.64 85.09
96.42 / 95.84
95.28
99.60 99.57
95.46
Task-specific
85.71 / 79.39
94.74 / 92.54
97.64
99.26 99.34
93.29
GPT-J
Union
72.61 64.89
89.68 / 87.76
95.56
99.82 99.83
88.21
Intersection
91.50 18.63
92.96 / 91.34
94.96
98.62 98.79
85.22
Varying the number of heads to be steered. Figures 3a and 3b illustrate the performance of PASTA when steering different number of heads on two tasks. The results suggest that as more heads are included for steering, the model follows the user even more closely, achieving higher efficacy (JSON Format Acc. and Pron. Change Acc.). However, at some point, this it results in a decrease in the metrics reflecting the generation quality (JSON Pred. Acc and Fluency). Thus, there is trade-off between emphasizing efficacy and generation quality. Overemphasizing can lead the model to focus solely on satisfying the user requirements and ignore the other parts. Therefore, we recommend applying PASTA to moderate number of heads (typically 50 to 150), striking balance between the efficacy and generation quality.
Varying the scaling coefficient a. Figure 3c presents the performance of PASTA on two tasks when we change the scaling coefficient The results indicate that PASTA is fairly robust to this
8

---

Published as a conference paper at ICLR 2024
100
100
80
Hormat Ace.
95
Acc.
Pinn. Ace.
05
ISON porma Ace.
Pin. Change Ace.
90
Fludley
90
x3
ISON Format Apc.
20
85
Pren. Chinge Ace.
40
JSON Formal Ace.
ISON Pred. Acc.
Fluescy
Pron. Change Acc.
Heedss
2
10
55 #lledd
56
108
294
0.002 10 10
(a) JSON Format
(b) Prons. Change
(c) Varying a
Figure 3: The performance of applying PASTA LLAMA-7B JSON Formating and Pronouns Changing tasks when varying the number steered heads [H] (3a,3b); and changing scaling coefficient a (3c).
hyperparameter; in practice, we fix it 0.01. Notice that setting to zero should be avoided, this leads to the complete removal of other crucial contexts at the steered heads, resulting in performance degeneration.
6 RELATED WORK
The primary method for controlling LLMs has been through prompting, often yielding impressive improvements in performance (Brown et al., 2020b; Liu et al., 2021; Wei et al., 2022) and spurring line of work aiming to make prompting easier, e.g. (Strobelt et al., 2022; Bach et al., 2022; Shin et al., 2020; Deng et al., 2022; Singh et al., 2023b). However, LLMs remain extremely sensitive to nuances in prompts (Webson & Pavlick, 2021; Lu et al., 2021); PASTA complements these approaches by making it easier for user to specify a prompt in difficult scenarios.
Another line of work aims to make LLMs more amenable prompting by modifying them during training. Most prominent among these approaches are instruction finetuning (Wei et al., 2021; Chung al., 2022), Reinforcement Learning from Human Feedback (Ziegler et al., 2019; Ouyang al., 2022), and other related methods, e.g. (Lee et al., 2023). There are also a few methods for directly specifying which parts on an input are important during training, e.g. (Ross et al., 2017; Rieger et al., 2019; Schramowski et al., 2020; Krishna et al., 2023). PASTA can be used in addition to these approaches to improve some aspects of model steerability (e.g. instruction following).
PASTA is related to variety of methods for adapting to new tasks, including LoRA (Hu al., 2021a), AdaLoRA (Zhang et al., 2023), QLoRA (Dettmers et al., 2023), and TOAST (Shi et al., 2023b). variety of research on model editing, e.g. ROME (Meng et al., 2022a),
PASTA is also related
MEMIT (Meng et al., 2022b), MEND (Mitchell et al., 2022), and REMEDI (Hernandez et al., 2023). Unlike these works, PASTA preserves an LLMs ability to transfer to new tasks using prompts and human-selected info, rather than using labeled examples.
Finally, PASTA is also motivated by works which have aimed to mechanistically understand attention scores (Zou et al., 2023). e.g. by studying them through feature importance (Jain & Wallace, 2019; Wiegreffe & Pinter, 2019; Deb et al., 2023), probing (Conneau et al., 2018; Liu & Avei, 2019), visualization (Karpathy et al., 2015; Olah et al., 2017), localizing knowledge (Meng et al., 2022a; Dai al., 2021), categorizing directions in representation space (Kim et al., 2017; Schwettmann et al., 2021), or natural-language explanations (Bills et al., 2023; Singh et al., 2023a).
7 CONCLUSION
In this study, we propose PASTA, novel approach aimed at enabling LLMs move beyond the limitations of plain text and effectively perceive user guidance embodied as highlighted parts of prompts. By making precise adjustments attention scores selected heads, PASTA directs the model's ss focus to the relevant context, mirroring the way humans benefit from textual cues. Unlike traditional fine-tuning methods, PASTA is applied at inference time and requires neither parameter updates nor gradient computation; PASTA requires only selecting which attention heads apply the re- re-weighting to, a one-time profiling operation for LLM. Experimental results show that PASTA significantly improve model performance variety of tasks. In the future, plan integrate PASTA with various other methods, such as few-shot in-context learning, aiming to highlight effective examples enhance its stability.

---

Published as a conference paper at ICLR 2024
REFERENCES
Stephen H Bach, Victor Sanh, Zheng-Xin Yong, Albert Webson, Colin Raffel, Nihal V Nayak, Abheesht Sharma, Taewoon Kim, M Saiful Bari, Thibault Fevry, et al. Promptsource: An integrated development environment and repository for natural language prompts. arXiv preprint (rXiv:222..01299 2022.
Steven Bills, Nick Cammarata, Dan Mossing, Henk Tillman, Leo Gao, Gabriel Goh, Ilya Sutskever, Jan Leike, Jeff Wu, and William Saunders. Language models can explain neurons in language models. URL https://oeeniipublcc. blob. core. windows. nethheuroon-cplaaiee/ppppeecceee 2023), 2023.
huml.(Date accessed: 14.05.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss,, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse. Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christo- pher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models few-shot learners. In H. Larochelle, M. Ranzato, Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020a. URL https: /ttttps::ppooooo /proeed ngs. www.rre..cppaappppppeeeeeee wttpcccppppppppeeeeee lss/ pptp/pppp//22///11eee.. 1457c0d6bEcb4966 wwwpaaaacaaaaaaaaaaabee
pdf.
Tom Brown, Benjamin Mann. Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances neural information processing systems, 13:1877-1901, 2020b.
Hyung Won Chung. Le Hou, Shayne Longpre, Barret Zoph. Yi Tay. William Fedus, Eric Li, Xuezhi Wang. Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finetue finetuned language models. arXiv preprint arXiv:2210. 1416, 2022.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? an analysis of BERT's attention. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Inter- preting Neural Networks for NLP. pp. 376-286. Florence, Italy, August 2019. Association for Computational Linguistics. doi: 11118653/1/W1194930 URL https: / a aclant org /W1 9828.
Alexis Conneau, German Kruszewski, Guillaume Lample, Loic Barrault, and Marco Baroni. What you can cram into single vector: Probing sentence embeddings for linguistic properties. arXiv preprint arXiv: 1805000000 2018.
Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. Knowledge neurons in pretrained transformers. arXiv preprint (xii::104400666 2021.
Maria De-Arteaga, Alexey Romanov, Hanna Wallach, Jennifer Chayes, Christian Borgs, Alexandra Choulde- chova, Sahin Geyik, Krishnaram Kenthapadi, and Adam Tauman Kalai, Bias in bios: A case study of semantic representation bias in high-stakes setting. In proceedings of the Conference on Fairness, Accountability, and Transparency, pp. 120- 1200-18.200 2019.
Mayukh Deb, Björ Deiseroth, Samuel Weinbach, Patrick Schramowski, and Kristian Kersting. Atman: Understanding transformer predictions through memory efficient attention manipulation. arXiv preprint arXiv:2301.08// 2023.
Mingkai Deng. Jianyu Wang, Cheng-Ping Hsich, Yihan Wang, Han Guo, Tianmin Shu, Meng Song, Eric P Xing, and Zhiting Hu. Rlprompt: Optimizing discrete text prompts with reinforcement learning. arXiv preprint 2205 2022.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized Ilms, 2023.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng. Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, Lei Li, and Zhifang Sui. A survey on in-context learning, 2023.
Evan Hernandez, Belinda Z. Li, and Jacob Andreas. Inspecting and editing knowledge representations in language models, 2023.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li. Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, arXiv preprint a0Xiv:2106.0940 09685, 202
J. Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint abs:2106 (ate:2206.09683 2021b.
10

---

Published as a conference paper at ICLR 2024
Sarthak Jain and Byron C Wallace. Attention not explanation. arXiv preprint 1rXi::900210186. 2019.
Andrej Karpathy, Justin Johnson, and Li Fei-Fei. Visualizing and understanding recurrent networks. arXiv preprint httl:::5000000000 2015.
Been Kim, Martin Wattenberg, Justin Gilmer, Carrie Cai, James Wexler, Fermanda Viegas, and Rory Sayres. Interpretability beyond feature attribution: Quantitative lesting with concept activation vectors (tcav). arXiv preprint 1r1i 1279, 2017.
Satyapriya Krishna, Jiaqi Ma, Dylan Slack, Asma Ghandeharioun, Sameer Singh, and Himabindu Lakkaraju. Post hoc explanations of language models can improve language models. arXiv preprint arXiv:2305. 1426, 2023.
Harrison Lec, Samrat Phatale, Hassan Mansoor, Kellic Lu, Thomas Mesnard, Colton Bishop. Victor Carbune, and Abhinav Rastogi. Rlaif: Scaling reinforeement learning from human feedback with ai feedback. arXiv preprint 20Xi02230.00026, 2023.
Frederick Liu and Besim Avei. Incorporating priors with feature attribution text classification. arXiv preprint arXiv:1906. 08286, 2019.
Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, Percy Liang. Lost the middle: How language models use long contexts, 2023.
Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, Graham Neubig. Pre-train, prompt. and predict: A systematic survey of prompting methods in natural language processing, arXiv preprint arXiv:2107.13586 2021.
Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. arXiv preprint arXiv:2104.0877 2021.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems, 35:17359- -13722, 2022a.
Kevin Meng, Arna Sen Sharma, Alex Andonian, Yonatan Belinkov, and David Bau. Mass-editing memory in transformer. arXiv preprint 1rXiv22210.07229, 2022h.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? In Н. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: httpss////p.ccceeng.g...ee.ee 2c601 (aweee2299ccb222676edeeeee
Eric Mitchell, Charles Lin, Antoine Bosselut, Chelsea Finn, and Christopher D. Manning. Fast model editing scale, 2022.
Chris Olah, Alexander Mordvintsev, and Ludwig Schubert. Feature visualization. Distill, 2(11):e7, 2017.
OpenAl. Gpt-4 technical report, 2023.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedhack. Advances Neural Information Processing Systems, 35:27730-27744, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, regh-performncc deep learning library. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer. Florence d' Alché-Bue, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 8024-8035. 2019.
Laria Reynolds and Kyle McDonell. Prompt programming for large language models: Beyond the few-shot paradigm, 2021.
Laura Rieger, Chandan Singh, W James Murdoch, and Bin Yu. Interpretations are useful: penalizing explanations align neural networks with prior knowledge, arXiv preprint arXiv: 19.999999.8.. 2019.
11

---

Published as a conference paper at ICLR 2024
Andrew Slavin Ross, Michael C Hughes, and Finale Doshi-Velez. Right for the right reasons: Training differentiable models by constraining their explanations. arXiv preprint 173.:1773..3777, 2017.
Patrick Schramowski, Wolfgang Stammer, Stefano Teso Anna Brugger, Franziska Herbert, Xiaoting Shao, Hans- MeonsGGerg Luigs, Anne-Katrin Mahlein, and Kristian Kersting. Making deep neural networks right for the right scientific reasons by interacting with their explanations. Nature Machine Intelligence, 2(8):476-44) 2020.
Sarah Schwettmann, Evan Hemandez, David Bau, Samuel Klein, Jacob Andreas, and Antonio Torralba. Toward visual concept vocabulary for gan latent space. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 6804-6812, 2021.
Tianhao Shen, Renren Jin, Yufei Huang. Chuang Liu, Weilong Dong, Zishan Guo, Xinwei Wu, Yan Liu, and Deyi Xiong. Large language model alignment: A survey, 2023.
Baifeng Shi, Siyu Gai. Trevor Darrell, and Xin Wang. Toast: Transfer learning via attention steering. arXiv preprint abw:2305. 155442, 2023a.
Baifeng Shi, Siyu Gai, Trevor Darrell, and Xin Wang. Refocusing key transfer learning. arXiv preprint arXiv:2305. 2023b.
Taylor Shin, Yasaman Razeghi, Robert L. Logan IV, Eric Wallace, and Sameer Singh. Autoprompt: Eliciting knowledge from language models with automatically generated prompts. arXiv preprint arXiv:2010.15980 2020.
Chandan Singh, Aliyah R Hsu, Richard Antonello, Shailee Jain, Alexander G Huth, Bin Yu, and Jianfeng Gao. Ex- plaining black box text modules in natural language with language models. arXiv preprint arXiv:2305.09863. 2023a.
Chandan Singh, John X. Morris, Jyoti Aneja, Alexander M. Rush, and Jianfeng Gao. Explaining patterns in data with language models via interpretable autoprompting, 2023b.
Nisan Stiemnon, Long Ouyang, Jeff Wu, Daniel M. Ziegler, Ryan J. Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul Christiano. Learning to summarize from human feedback, arXiv preprint 205:2000.013255 2020.
Hendrik Strobelt, Albert Webson, Victor Sanh, Benjamin Hoover, Johanna Beyer, Hanspeter Pfister, and Alexander M. Rush. Interactive and visual prompt engineering for ad-hoc task adaptation with large language models, 2022.
Ian Tenney, Dipanjan Das, and Ellie Pavlick. BERT rediscovers the classical NLP pipeline. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4593-4601, Florence, Italy, July 2019. Association for Computational Linguisties. doi: 10.18653/v1/P19-1444 1452. URL https: //aclantee anchont hoy.yyoooooooooooooooooooooooo.. , Arg//199 9- 1452.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasminc Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, F. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL attps::/proceed:e http://pooeeeen...eeee www.iis............. cc.cccpppaa.. co/paper. files/
Elena Voita, David Talbot, Fedor Moiseev, Rico Semnrich, and Ivan Titov. Analyzing multi-head self self-attention: Specialized heads do the heavy lifting, the rest can be pruned, July 2019. URL URL.attss//accaanhooooo
org/p19-155 -1580.
Ben Wang and Aran Komatsuzaki, GPT-J-6B: A Billion Parameter Autoregressive Language Model. https: //github. com/ kingofionni me trans f oormer jax. May 2021.
Albert Webson and Ellie Paylick. Do prompt-based models really understand the meaning of their prompts? arXiv preprint 1XX::200900124: 2021.
Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoe V Le.. Finetuned language models are zero-shot learners. arXiv preprint arX1::2109.01652 2021.
12

---

Published as a conference paper at ICLR 2024
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, al. Chain-of-thoughh prompting clicits reasoning large language models. Advances Neural Information Processing Systems, 35:24824-2483 2022.
Sarah Wiegreffe and Yuval Pinter. Attention not not explanation. arXiv preprint arXiv:1908.04022 2019.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Huggingface transformers: State-of art natural language processing. arXiv preprint 1rXiv:191000377,, 2019.
Zhewei Yao, Reza Yazdani Aminabadi, Olatunji Ruwase, Samyam Rajbhandari, Xiaoxia Wu, Ammar Ahmad Awan, Jeff Rasley, Minjia Zhang, Conglong Li, Connor Holmes, Zhongzhu Zhou, Michael Wyatt, Molly Smith, L A Kurilenko, Heyang Qin, Masahiro Tanaka, Shuai Che, Shuaiwen Leon Song, and Yuxiong He. Deepspeed-cher Easy, fast and affordable rlhf training of chatgpt-like models at all scales, arXiv preprint abs:2308.01 0330000000.00. 2023.
Qingru Zhang, Minshuo Chen, Alexander Bukharin, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adaptive budget allocation for parameter-efficien fine-tuning. In The Eleventh International Conference on Learning Representations, 2023. URL https: : http//peeeeeeeeeee/eee/ee/eeevve net/ www.ooiiuoleeeooeoo wet.nniddddl@eeeeeu jiY.
Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffre Irving. Fine-tuning language models from human preferences. arXiv preprint 2019.
Andy Zou, Long Phan, Sarah Chen, James Campbell, Phillip Guo, Richard Ren, Alexander Pan, Xuwang Yin, Mantas Mazeika, Ann-Kathrin Dombrowski, Shashwat Goel, Nathaniel Li, Michael J. Byun, Zifan Wang. Alex Mallen, Steven Basart, Sanmi Kayejo, Dawn Song. Matt Fredrikson, J. Zico Kolter, and Dan Hendrycks. Representation engineering: A top-down approach to ai transparency, 2023.
13

---

Published as a conference paper at ICLR 2024
APPENDIX
A EXPERIMENTAL DETAILS
We implement all algorithms using PyTorch (Paszke et al.. 2019) and buggingface (Wolf et al., 2019) and run experiments on NVIDIA V100 GPUs and NVIDIA A6000 GPUs. Table 6 provides detailed statisties of datasets in our experiments.
Table 6: Statistics of datasets.
Task
Train
Valid
Test
CounterFact
1000
1000
5000
BiasBios
1000
1000
5000
JSON Formatting
1000
1000
5000
Pronouns Changing
1000
1000
5000
A.1 DETAILED PROMPT TRMPLATES EACH TASK
For each task, the prompt templates in our results are as follows:
JSON Formatting:
(Original) {context). Answer the occupation of [person} and generate the answer as json format. Here is an example: {"name": "occupation". .). Now generate the answer.
(Shortened one in Section 5.2) {context). Answer the occupation of {person} and generate the answer as json format.
(Rephrased one in Section 5.2) Answer the occupation of [person) and generate the answer as json format. Here is an example: {"name": ""ocuupation" ": {context). Now generate the answer.
Pronouns Changing:
(Originaly: {context). For the aforementioned text, substitute 'she' and 'he" with 'they' and generate the occupation of {person) after changing pronouns.
(Shortened one in Section 5.2): {context). Change 'she" and "he" with 'they' and answer the occupation of {person) after replacing the pronouns
(Rephrased one in Section 5.2): {context). For the aforementioned descriptions, replace 'she' and 'he' with 'they' the aformentioned text and generate the new text after replacing the pronouns.
BiasBios: (context). {person} has the occupation of.
CounterFact:
Previously, fold fact). Currently, (new fact). {question)
A.2 THE EVALUATION DETAILS OF PASTA
Table 7 presents the number of heads to be steered by PASTA for LLAMA-7B and GPT-J-6B on every task.
Table 7: The number of heads to be steered by PASTA.
Task
LLAMA-7B
GPT-J-6B
JSON Formatting
53
153
Pronouns Changing
86
72
BiasBios
86
111
CounterFact
86
52
14

---

Published as a conference paper at ICLR 2024
B EXTENDED RESULTS
B.1 EXTENDED RESULTS WITH FLUENCY
In this section, include extended results, including fluency metrics. Fluency score is the average bi- gram and tri-gram entropy of generations, designed to be low for degenerated or repetitive texts (Meng al., 2022a). This metric be regarded as the reference metric of generation quality. Typically, the generations of language models are reliable as long as their fluency score is not too low. Here, we filter out any results receiving fluency score below 3.0. Table 8, 9 and 10 include all results and fluency evaluation.
Table Main results LLAMA-7B to demonstrate that PASTA can improve the model ability to (i) follow user instruction (JSON Format and Prons. Changing); (ii) interpret contextual information (BiasBios); (iii) resolving knowledge conflicts (CounterFact). For all scores, higher is better. The best results are in bold.
JSON Format
Prons. Changing
BiasBios Ace Flue.
CounterFact ES PS /Flue.
Method
F. Acc / P. Acc
Acc A.Ace Flue.
Zero-shot
60.00 54.94
71.84 / 66.28 6.10
87.36 3.98
58.50 52.03 4.96
t-marked
18.55 12.71
39.14/ 35.17 6.03
90.62 3.89
57.74 50.52 5.12
Prompting
" ""-marked
4.56 4.20
20.55 18.19 5.13
89.82 3.97
38.14/51.77 5.13
Few-shot
84.85 /3.58
59.06 / 55.27 5.95
88.79 4.19
87.45 / 49.82 5.68
Task-agnostic
88.16 49.08
83.65 81.31 4.62
/
98.82 / 99.03 4.78
PASTA
Multi-task
96.64 85.09
96.42 95.84/ 5.43
95.28
4.05
99.60 / 99.57 4.89
Table 9: Main results of GPT-J to demonstrate that PASTA improve the model ability (i) follow user instruction (JSON Format and Prons. Changing); (ii) interpret contextual information (BiasBios); (iii) resolving knowledge conflicts (CounterFact). For all scores, higher is better. The best results are bold.
JSON Format
Prons. Changing
BiasBios
CounterFact
Method
F. Acc / P. Ace
Ace / A.Acc / Flue.
Acc / Flue.
BS / PS /Flue.
Zero-shot
28.83 / 25.09
39.88 / /
72.76 /2.76/5.00 5.06
22.14/42.02// 5.01
markee
/4.11
41.25 1/25/337 37.57 /476
74.14/5.01
44.50 / 45.09 / 5,22
Prompting
""-marked
8.81 5.62
6.12/
78.64 4.96
45.54 1..4//41.0 41..// 5.16
Few-shot
84.15 72.65
35.77 32.08 / 6.46
72.98 4,82
68.34 / 38.23/ 5.67
Task-agnostic
46.68 / 34.71
91.62 / 88.60 3.00
80.84 / 4.92
99.54 / 99.57 5.11
PASTA
Multi-task
91.50 18.63
92.96 / 91.34 4.91
94.96 4.87
98.62 98.79/ / 5.11
Table 10: Varying head selection strategies between top top task-specific heads, union across multiple tasks, and intersection (the default used in PASTA).
JSON Format F. Acc / P. Acc
Prons. Changing
BiasBios Ace / Flue.
CounterFact
PASTA
Acc/ A.Ace Flue.
ES / PS /Flue.
GPT-J LLAMA
Task-specific
95.56 / 86.83
98.52 / 98.02 5.92
97.62 4.18
99.18 / 99.24 4,93
union
88.42 / 74.49
92.12/91.44 4.88
96.36 4.13
99.24 / 99.35 / 44553
intersection
96.64/ / 85,09
/95.849 5.43
95.28 4.05
99.60 / 99.57 / 4.89
Task-specific
85.71 / 79.39
94.74/92.54 55007
5.07
97.64 5.06
99.26 / 99.34 / 4.94
Union
72.61 / 64.89
89.68 / 87.76 3.92
95.56 5.02
99.82 / 99.83 5.03
Intersection
91.50/ / 18.63
92.96 /91.34 4.91
94.96 4.87
98.62 / 98.79 5.11
B.2 THE VARIANCE OF FEW-SHOT PERFORMANCE
Few-shot prompting sometimes leads to improvements in model performance. as explicitly providing the examples in additional demonstrations, However, drawback of few-shot prompting is its insta-
15

---

Published as a conference paper at ICLR 2024
bility, meaning its performance exhibits high variance across different samples in the demonstratio. In this section, we present the results to show that the performance of few-shot prompting displays high variance in temms of sampling different few-shot demonstrations.
Table 11: The few-shot performance (Ace. / A. Acc. / Fluency) on the Pronouns Changing task.
Few-shot examples
ILAMA-7B
GPT-J-6B
Demonstration I
84.87 /9.09/ / 4.74
43.82 40.36 6.43
Demonstration 2
57.24 / 53.98 6.22
40.68 37.86 / 6.44
Demonstration 3
57.08 53.22 6.02
33.13 29.21 6.48
Demonstration 4
52.26 48.30 6.42
25.47 20.89 6.44
Demonstration 5
43.86 / 40.78 6.43
11.90 / 6.63/6.51
B.3 MODEL PROFILING RESULTS
In this Section, provide more results of the performance of LLAMA-7B on all of tasks when steering: (i) all heads; (ii) entire layer; (iii) a individual head of a layer.
100
80
ACC
60
Zero shot
Steer
Steer entiree layer
20
Steer single head
12
24
Layer
Figure 4: The performance of LLAMA-7B on Pronouns Changing task when we steer (i) all heads (green); (ii) entrie layer (yellow); and (ii) individual head with a layer (blue violin plot). The performance varies dramatically across layers and across heads of layer.
70
PASTA
heads
60
layer
Steer single head
12
16
20
24
32
Layer
Figure 5: The performance of LLAMA-7B on BiasBios task when steer heads (green); (ii) entrie layer (yellow); and (iii) individual head with a layer (blue violin plot). The performance varies dramatically across layers and across heads of a layer.
100
80
Zerc-shet
Steer heads
40
Steer entire layer
single head
12
24
Layer
Figure 6: The performance of LLAMA-78 on CounterFact task when we steer (i) all heads (green); (ii) entrie layer (yellow); and (im) individual head with layer (blue violin plot). The performance varies dramatically across layers and across heads a layer.
16

---

Published as a conference paper at ICLR 2024
C RESULTS ON MORE MODELS
C.I LARGER MODEL SIZE
We conduct experiments with LLAMA-13B to further evaluate the effectiveness of PASTA across all tasks. The following table presents the performance comparison for LLAMA-13B.
Table 12: Results of LLAMA- 13B. For all scores, higher is better. The best results are in bold.
JSON Format F. Acc / P. Acc
Prons. Changing
BiasBios Acc
CounterFact
All
Method
Acc A.Acc
ES PS
Ave.
Zero-shot prompting
45.48 44316
65.03 / 60.90
85.80
47.86 44. 14
56.05
Few-shot prompting
39.80 / 3.56
82.33 / 80.71
88.38
90.63 65,49
64.41
PASTA (Multi-task)
98.74 89.88
97.56 / 96.78
95.34
99.38 99.30
96.71
C.2 STEERING INSTRUCTION-TUNEE MODELS WITHOUT DE-PROFILING
We further test PASTA's upplicability to Vicuma-7B-vl. which is instruction-turs tundd from LLAMA- 7B. We apply PASTA using attention heads selected from I.AMMA-B profiling (including multi-task and task-specifie heads). In this way, we evaluate if the heads selected from the base models are transferable to an instruction tuned model, thereby avoiding the re-profiling. The table below presents the performance of Vicuna across all tasks.
Table 13: Results of Vicuna-7B. For all scores, higher is hetter. The hest results are in bold.
JSON Format
Prons. Changing
BiasBios Acc
CounterFact
Method
F. Ace / P. Acc
Acc/ A.Acc
ES / PS
LLAMA-7B Zero-shot
60.00 / 54.94
71.84 66.28
87.36
58.50 / 52.03
LLAMA-7B PASTA(multitaskk)
96.64 / 85.09
96.42 95.84
95.28
99.60 / 99.557
Vicuna Zero-shot
65.41 / 61.78
95.74 94.74
90.74
61.10 / 52,46
Vicuna PASTA (ASTA(multi-task)
66.09 / 56.00
98.82/ 98.08
96.44
99.80/ / 99.80
Vicuna PASTA(tas--specifi))
90.54 / 86.53
98.62 98.04
97.42
99.82 99.74
The results demonstrate that the attention heads selected for LLAMA-7B effectively steer Vicuna-7B, indicating that re- re-profiling is not necessary for instruction- models. Notably, when steering task- specific heads selected from LLAMA profiling, PASTA significantly enhances Vicuna's performance across all tasks. This evidence shows that PASTA can complement instruction tuning without necessitating re-profiling.
C.3 ABLATION ABOUT THE NUMBER OF EXAMPLES FOR PROFILING
the robustness head performance ranking sample variance allows us further reduce the sample size for profiling (e.g., |D] 200). The table below presents the PASTA performance the JSON Formatting task when re- re-profiling with (D) 200 samples. We can see PASTA still achieves superior performance when profiling with much fewer examples.
Table 14: The performance of PASTA with different sample size |D] of model profiling.
Model
Sample size
ISON Format Acc
JSON Pred. Acc
LLAMA- 7B Zero-shot
N/A
60.00
54.94
LLAMA-7B W. PASTA
1000
95.56
86.83
LLAMA-7B W. PASTA
200
95.04
84.17
Vicuna-7B Zero-shot
N/A
65.41
61.78
Vicuna-7B w. PASTA
200
98.56
94.56
17

---

Published as a conference paper at ICLR 2024
D EXAMPLES OF PASTA GENERATIONS
Table 15: Generation examples of LLAMA-7B on JSON Formatting.
Prompt
PASTA Generation
Winnie is an American photographer living in New York. Specialized fashion photography and portrait, she applies her talent both humans and animals. ln this scries of pictures, It's all about personification those dogs in spotlight The whole thing gives a great result. More pictures & infos on her portfolio. (instmaction)
"mmmm": "Winnie", "necupation" \nn"nmee: "Winnie", "occupation",
"photographer" "photographer"
Jon is software engineer on the Core Processing team at comScore, online audience measurement and analyties company. Hc prefers to do most of his coding in Pig. Before working at comScore be wrote soft- optimice supply chains and to allovate lixed income securities. (instruction)
"Jon", "occupation" neer") ['nn" [nnamme"" "Jon", "occupation": "Software Engi- nser"). "..")..
("name":
"Software Fini--
Juris an associate professor the Faculty of Information, University Toronto. He co author of Readers Advisory Service in North Ameri- can Public Libraries, 1870-2005 (MeFarland, 2007); Evolution Library and Museum Partnerships: Historical Antecedentss Conterppo- rary Manifestations,
("mame: "Juris", "occupation". "Associate Professor") (n{ jn("name; "Juris" "occupation" "Associate Professor" J...
Future Directions (1 iimurriss Unlimited, 2004); and Reading and the Reference Librarian: The Importance to Library Service of Staff Reading Habits (McFarland. 2004). ( [instruction]
Gregory is an educational psychologist whose major concern is with how people learn, ie skills and knowledge development, expecially in educattional settings. His emphasis is strong research (statistical) foundation crossing arcas such as social learning theory, information processing, cognitive approach to emotional lactors. (instruction)
("name"; "Gregory", "occupation": "Educational Psycholo- gist") } (nn"mmm"" "Grregury, "occupation", : "TEducatioaal Psy. chologist") J...
Asif a Solutions Architect with Amazon Web Services. He provides tochnical guidance, design advice and thought leadership some the largesi and successful AWS customers and partners on the planet. His deepeat expertise spans application architecture, containers, devops, security, machine learning and SmiS business applications. Over the last 12 years, he's brought an intense customer facus to challenging and deeply technical roles in multiple industries. He has number of patents and has succesfully led product development, architecture and customer engagements. { {instruction] }
{ ("name": "Asif", "occupation". "Solutions Archi- teet")) "aa"""a""mme" chitect" J...
"Asif", "occupation" "Solutions
Grahan is an Associate Professur in Construction Management at the University Neweastle, Australia, where he end of 1999, He is currently Depoty Head of the School of Architecture and Built Environmeat, and was previously Head of Discipline, Building. During this time he has undertaken various administrative roles School, Faculty, and University level, and is currently member of Academic Sonate. (instruction)
{ ("mume: "Graham"", "uscapation": "Assaciate Professor in Con- struction Management") ") \a ("name": "'Graham", "occupation" "Associate Professor in Construction Management") )
appointed at the
Katherine is an Assistant Professor in the Department of Pathology. She received her MD where she also completed her APICP resädency training in pathology completed her fellowship
("name": "Katherine", "oceupation": "Assistant Professor in the Department Pathology" : "answer": "Katherine an Assistant Professer in the Department of Pahhloog. She reveived her MD she of Medicine
University of Colorado School Medicine
and a fellowship in Cytogenetics. Molecular Genetic Pathology in the Department of Pathology in 2009. She is a Medical Director in Cytogenctics, Molecular Cytogeneties, and Molecular Oncology ARUP Laboratories. (instruction)
University Colorado also completed her APICP residency training in pathology and fellowship in Cytogeneties."
Peter an immiggaation attorny with least three links to the Mexi- can government. addition to heading the Center for Human Rights and Constitutional Law, he also operates vocesunidas.or which until recently was called "Mexico Project". It's since undergone an inter- esling change. (instruction]
[ ("mume": "Peter", "occupation": "Immiigration Attorney "] 1/nn object. ("name": "Peter", "occupation"
The answer "an" ""ammee" "Poter", "occupation" "Immigration Attorney") Ao "Immigration Attorney") Y...
Delvin freelance joumaaiis.. was of Chippewn's Oil and Gas Specialist, He frequent contributor to the tribe's newspaper. The Turtle Mountain Times, as well as indianz. com and Indian Country Today Media Network, both national news outlets for Indian Country. (instraction)
Turtle Mountain Band
""name": "Delvin", "occupation" "Freelance Journalist)) ) (" ("name"; "Delvin", "occupation": "Freclance Journalist" "age": 25H ("name") "Delvin", "occupation". "Freclance Journal- iss". "age": 25, "gender"; "made") ("name": "Delvin", "oc- cupation", "Freelance Journaiss",, "age": 25, "gender": "male". "hair": "black", "eyes": "brown")...
Robert is a rock photographer based in Los Angeles. He started his ca- reer staff photographer for the legendary Creem Magazine. He pho- tographed the Stouges al some of their earliest shows and has maintained a close working relationship with group. His woek appeared Playboy, The New York Times, Los Angeles Times, Rolling Stone, and Mojo, as well as on over 200 albums, including the Stooges' most recent album, The Weirdness. He oversaw the 2001 return of CREEM maga- zine, as well as the 2007 hankooeer anthology CREEM: America201 9 Only Rock 2018a2019 Magazine. (instruction)
{ ("nume": "Robert Matheu", "occupation": "Rock Photogra- pher") An "ud"umm"" "Robert Matheu", "occupation", "Rock Pho- tographer"). )....
18

---

Published as a conference paper at ICLR 2024
Table 16: Generation examples of LLAMA-77 on Pronouns Changing.
Prompt
PASTA Generation
Flue./oons.
Jan software engineer Core Processing team comScore, an online audience measure- ment and analyties company. He prefers to do most of his coding Pig. Before working comScore he wrote software to uptimize supply thains and allocate fixed income securities. (instruction) }
software engineer on the Cure Processing learn at comScore, an online audience measurement and analyties company. They prefer to most of their coding in Pig. Before working comScore they wrote software optimize supply chains and to allocate lixed income securities. ThTh ahove text is a good example of a gender neutral sentence. The following sentence: \ nlon is a software engineer on the Core Processing team at comScore, an online audience measurement and analytics company. They prefer to do most of their coding in Pig.
5.95/19
good example a gender-neutral
Fred is a photographer who also runs painting and wallpapering business in Brooklyn. He has exhibited his work in group shows in Brooklyn and Manhattan Hope & Auchor restaurant, Object Image Gallery, Foley Gallery and BAX (Brooklyn Arts Exchange). He primarily self-taught studied at the New School and the International Center of Photography, (instruction):
Fred photographer who also runs a painting and wallpapering husiness in Brooklyn. They have exhibited their work in group shows in Brooklyn and Manhattan at Hope & Anchor restaurant. Objoct Image Gallety, Foley Gallery and BAX (Brooklyn Arts Exchange). They are primarily self taught and studied at the New School and International Center of Pbotography, /ntTe text now gender neutral. / nhThe text is now gender neutral. The pronouns have been changed to 'they" 'them' and the occupation has heen changed photographer"
6.39 6.39/28.16
a journalist who writes on travel par- enting for The New York Times Magazine, Vogue, Conde Nasi Traveler and many other publications. Adrienne Hartman is illustratee living New York. Her work has been published by Gournet, Moath Club, and other publi-
Ann a joumalist writes on travel parenting for The New York Times Magazine, Vogue, Conde Nast Traveler and many other publications. They are a journalist who writes on travel and pareming for The New York Times Magazine, Vogue, Conde Nast Traveler and many other publications. AnAnn is journalist who writes on travel and parenting for The Nast Traveler and many other publications.
2.99/20.91
Entrece. Book cations and books. (instruction)
York Times Magazine, Vogue, Conde
Justine photographer. Photo District News named her one of their top 30 new and emerging phorographers to watch in 2011. She also rocently received Queens Community Arts Fund grant from the Queens Council on the Arts, She has held many artist residencies, Lower Manhattan Cultural Council. Her work has exhibited at many venues, including Michael Muzzeo Gallery and El Museo del Barrio in New York. She earned a Master of Fine Arts from San Franciscu Art Institute, ( (instruction]
Justine is photographer. They are photographer. cettrappher//Justtinee AuJustine photographer. They are a photographer. They are photogra- pher.' beler'alusse alustine a photographer. They are photographer. They are photographer. They are photographer. AnJustine a photographer. They are photographer. They are a photographer They are a photeg- rapher. They photographer. AnJustine is photographer. They are a photographer. They are a photographer They are a photographer. They are a photographer. They are a photographer
2.89 14.18
recently at the
Courtncy is a freelance journalist living who pre- viously lived Buenos Aires. She writes for True/Slant, Packet Lint, Flavorwire, Relix, GBQ. Wander-Argeniing & Forhes She loves rohots, tor- rontes, yogu and live music. She was named one Rusiness Insider, Most Important Women Un- transatlantic der 30' and founded audience.io, company whose aim is to help New York City and London based start ups collaborate internationall. She even has her own Wikipedia page.
Courtney is freelance journalist living who previously lived Buenos Aires. They write for True/Slant, Pocket- Lint, Flavorwire. Relix, GBQ, Wander- Argentina & Forbes. They love robois, tomontes. yoga and live music. They were named one of Business Insider, Most Important Women Under 30' and founded audience.io, transatlantic company whose aim help New York City and London-hased start- ups collaborate intemationally. They even have their own Wikipedia.
6.33 / 12.17
Gavin is Sydney photographer focused on people and places. He specialises in commercial and edi- torial photographt and is regularly commissioned to work on variety of corporate, editorial, archi- tectural, industrial, advertising and annual report projocts. He is highly vessatile portrait and loca- tion photographer. [instruction]
Gavin Sydney photographer focused on people and places. They spocialise in commercial and editorial photography and are regularly commissioned work a variety corporate, editorial, architec- tural, industrial, advertising and annual report projects. They are a highly versatile postrait and location photographer.
5.07 13.72
Eric professional photographer. Scott Allen Perryis an actor, an independent filmmuker, and a writer. He the writer and director "Bobbi's Bedroom," "MC Outdoorz," and the decumentary "The Ouddorrsenn. They both live in Los Angeles. Adam Mock is the president and creative director Medallion Modia Group. He lives in Mont- gomery, Illinois. They are the pholographer and authors of "Mime Very Own Bank." (instruction)
Eric is professional pholographer. They are the photographer and authors of "Mime Very Own Book."((Theei occupation a profes- sional photographer /nThii occupation is a professional photogra- pher. They are the photographer and authors of "Mime Very Own Book."
4.53/17.31
Christopher freelance journalist who grew up in Michigan. He has a journalism degree from Cen- tral Michigan University and spent six years writ- ing for the Upper Peninsula's is largest daily newspa- per. an amateur adventurer and explorer and knows the socret location of Marquette County's Big Boy graveyard He designs newspapers for living.
Christopher is freelance journalist who grew up Michigan. They have a journalism degree from Central Michigan University and spent years writing for the Upper Peninsula's largest daily newspaper. They are an amatcur adventurer and explorer and know the secret location Marquette County's 's Big Boy graveyard. They design newspapers for living.
5.58 / 13.85
19

---

Published as a conference paper at ICLR 2024
E ADDITIONAL EVALUATION METRICS
We understand the importance of preserving generative flueney and quality while enhancing task- specific performance with PASTA. To ensure this, we employ two metrics to evaluate the quality of PASTA generations across three natural language generation tasks (Prons. Changing, BiasBios, and CounterFact).
Fluency Evaluation (Meng et al., 2022a): As mentioned in Section 5, we assess the fluency of all generations (the average bigram and trigram entropy of generations), and exclude results with fluency repetitive generations from consideration.
below 3.0. This step effectively eliminates degenerated or
Consistency Metric: We employ an additional consistency metric (introduced by Hernandez et al. (2023)), which measures the average tf-idf similarity between the generated text and reference texts full dataset. This metric helps us measure how well the generated text aligns with overall contextual inputs in terms of content and style (higher is better).
Table 16 presents examples of LLAMA-7B generation with PASTA and their fluency and consistency scores on the Pronouns changing task. We can see that repetitive or meaningless generations receive low fluency (below 3.0) and consistency (below 8.0). The generations with high fluency (around 4.5) and consistency (above 13) meaningful and readable. The following table presents the average fluency and consistency evaluation across the mentioned tasks:
Table 17: Results of fluency and consistency evaluation LLAMA-7B.
Prons. Changing
BiasBios Acc Cons. / Flue.
CounterFact
Method
Ace Cons. / Flue.
ES /S// Cons. Flue.
Zero-shot
71.84 22.29 / 6.10
87.36 13.02 / 3.98
58.50/ 52.03 / 11.64 / 4.96
PASTA
92.30 22.33 6.07
95.28 14.25 / 99.60 99.57 19.29 4.89
The results show that PASTA achieves comparable consistency and fluency scores to zero-shot prompting. This indicates that PASTA effectively maintains the generative quality and fluency while significantly improving the task efficacy.
20


Reviews:
[{"review_id": "OSg7YqDsx0", "reviewer": "Reviewer_ppCr", "summary": "This paper proposes a novel attention steering approach, PASTA, to make the model focus on a subset of the prompt/instruction specified by the user at inference time for large language models. PASTA re-weights the attention scores by reducing the scores of non-highlighted tokens by a factor of $\\alpha \\sim 0.01$. After normalization, the attention scores of the highlighted tokens are increased. The modified attention matrix is used for computing the projections of the tokens, which are further passed to the feed-forward layers and eventually the next transformer layer. The paper also proposes a multi-task model profiling algorithm to select a subset of the heads in each layer of the model. Empirically, this leads to performance gains over simple attention head-selection heuristics or using all the attention heads. The paper conducts experiments on four synthetic tasks - JSON Formatting, Pronouns changing, Counterfact, and BiasBios using the LLama-7B and GPT-J models. The proposed approach achieves significant improvements over zero-shot, *-marked, \"\"-marked, and few-shot prompting methods.", "questions": "I have asked most of my questions in the weakness section, but here are some followup questions that I am interested in:\n\n- Beam search sampling and top-k sampling are much better compared to greedy. The baseline performance would definitely improve using these generation methods. How does PASTA compare to the baselines using these sampling approaches?\n- *-marked and \"\" - marked have very low performance. Have the authors tried better markers like <mark> </mark> or <emphasize> </emphasize>. I believe these markers are much more frequent in the training set, and might improve performance.\n\nAlso, there are a couple of formatting issues:\n- Appendix section A is empty in the submission. Please remove that in the updated version.\n- Citations and links are coloured. They should be black for ICLR.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well written and fairly easy to follow.\n- The proposed approach is novel and intuitive as increasing the attention scores of highlighted tokens will lead to more contribution in the output projection of tokens during generation.\n- The multi-task model profiling to select a certain subset of attention head improving performance on evaluation tasks leads to better generalization even on unseen tasks.\n- The empirical results on the four tasks show significant improvements over the standard zero-shot and few-shot prompting baselines.", "weaknesses": "- The paper mentions that the proposed approach, PASTA, does not need access to model weights, but I believe that is not correct? Once the attention scores have been modified, you would need access to $W_{v_{h}}$ to compute $H^{(l, h)}$ in equation 1. Furthermore, you'll need the weight matrices of the two feedforward layers $W_{ffn1}$ and $W_{ffn2}$ (assuming non-gated activation) to compute the output of a given transformer block layer, which can be passed to the next layer. Please correct me if I am wrong, but I think that the next layer's output cannot be computed without access to these weights.\n- The results are on those tasks only where the highlighted instruction is at the end of the prompt (except for the rephrased JSON Formatting task). What is the general effect of having highlighted instruction at the start or some other place in the prompt? This would give an idea if the proposed approach is general enough.\n- Not a weakness, but there's a typo in footnote 3: access $\\textbf{to}$ model weights ...", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ONRV7t7yXJ", "reviewer": "Reviewer_s9kK", "summary": "This paper introduces PASTA (Post-hoc Attention STeering Approach), which addresses the need to guide the attention of large language models (LLMs) towards user-specified information. Unlike existing methods that are limited to plain text processing, PASTA enables LLMs to read text with user-specified emphasis marks, mimicking the subtleties of text style in human-written articles. PASTA achieves this by identifying a select group of attention heads and applying precise attention reweighting, directing the model's focus towards user-specified parts of the text. The method is applied at inference time and does not require any changes to the model parameters. Experimental results demonstrate that PASTA significantly improves an LLM's ability to follow user instructions and integrate new knowledge from user inputs. The performance improvement is substantial across various tasks, with an average accuracy boost of 22% observed for LLAMA-7B.", "questions": "1.\tOne thing I am really curious about is whether we have any preferences layer-wise. Since different layers are known to possess different functions. Which layer is more important for attention steering? By Figure 2 it seems that it’s approximately random, is there any explanation for this?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1.\tThe authors propose a simple yet effective method called post-hoc attention steering. The method improves performance in multi-task and task-agnostic settings by forcing the model to focus on several important positions.\n\n2.\tThe structure of this paper is really easy to follow and understand.\n\n3.\tThe experiments are pretty comprehensive from my perspective, all the claims are properly proved.", "weaknesses": "1.\tExperiments on larger models should be conducted to ensure the performance across different models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iHTJ8dfnIq", "reviewer": "Reviewer_pBs7", "summary": "This paper proposes a post-hoc method, called PASTA, that allows users to steer the model's attention towards user-specified information by highlighting relevant parts of the input. The idea is to select a small subset of attention heads and applies precise attention reweighting on them; that is, we select some attention heads and steer the corresponding attention weight to focuses more on the specified input spans. This proposed method is evaluated on four tasks with GPT-J-6B and LLaMA-7B. Experimental results show that PASTA achieves an average accuracy improvement.", "questions": "- Did you perform the SFT produce on the models in the experiment? \n- You should present more details in computing accuracy on these used datasets.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This paper is easy to read and the method is not difficult to implement. The idea of steering attention to enhance large language models' adherence to instructions proffers significant value. This paper also introduces a PASTA method, which can select a small subset of attention heads and applies precise attention reweighting on them during the inference process.", "weaknesses": "- This proposed method necessitates the predefined highlighted input spans. However, in the context of numerous tasks, defining such input spans to be a formidable challenge. One idea is whether an automatic extraction mechanism could be designed to improve this method. This could entail the development of a dedicated extraction model, or the crafting of specialized prompts to guide Large Language Models in executing this task more effectively. \n- The PASTA method does not necessitate alterations in model parameters, resulting in enhanced performance efficiency; however, a pertinent issue arises concerning the potential disruption of the established generative law from language modeling via manual attention steering. This makes a conjecture that such intervention could inadvertently impair other abilities of LLMs, such as generative fluency and reasoning.\n- The current experimental design appears to be limited in scope. For a more comprehensive evaluation, it would be beneficial to test the effectiveness of the proposed method on diverse large language models, including but not limited to, 7B-LoRA, 13B-LoRA, and 13B, across specific tasks.\n- There have been previous studies on improving the standard attention with the prior knowledge, such as syntax-based attention and multiscale attention. These should be considered baselines for comparison, either in related work or experiments. Additionally, although the four tasks previously considered do not adequately represent the breadth of generative tasks commonly utilized as benchmarks, they do not represent the prevailing benchmarks within generative tasks. It is crucial to evaluate the proposed methodology's efficacy by applying it to some commonly used generative tasks, including but not limited to, machine translation and question answering tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3kDI7QRqSI", "reviewer": "Reviewer_WBEq", "summary": "This paper introduces a method to improve the quality of text generations on specific tasks by specifying explicitly to which portions of the prompt the model should pay close attention. To this end, the authors introduce 4 instruction tasks which have been annotated explicitly in order to place emphasis on a particular section of the prompt.\n\nThe authors run a \"model profiling\" phase on a training subset of each of these tasks. This phase aims at detecting which attention heads should see their outputs amplified in order for the LLM to follow more closely the operator's instructions. At inference time, the post-softmax attention scores of the corresponding heads are boosted by a constant factor \\alpha.\n\nThe results show strong gains against raw LLMs in zero-shot and few-shot setups, and interesting robustness properties related to prompting.", "questions": "* Could you clarify how we should interpret results on json generation and pronoun change given the shortcomings of the evaluation code?\n\n* Do the findings transfer on readily-available instruction-tuned models?\n\n* In effect, the prompt the paper is operating on has two components: the input text and the emphasis. In section 5.2, you demonstrate that PASTA mitigates noise and poor wording choices in the input test. For a full comparison with zero-shot prompting, could you test the effect of poor choice of emphasis in the input text?\n\n* In the same section, what is the reason to compare against zero-shot and not a few-shot setup? It would seem that few shot prompting would be a fairer comparison against multi-task PASTA, as both assume that they have access to some training examples on the task at hand.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 1, "presentation": 3, "contribution": 2, "strengths": "* Originality: the proposed method for increasing attention weights for sections that have been manually emphasized is novel. The method is relatively lightweight, which makes it compelling to adapt very large models dynamically to new tasks.\n\n* Quality: the code provided with the paper is problematic, as it means that the proposed evaluations may not correlate with human judgement on these tasks. In addition, a more thorough comparison with off-the-shelf instruction-tuned models may be relevant. See Weaknesses section for more details.\n\n* Clarity: the paper is well-written and easy to follow. However, the definition of the main metrics could be updated to more closely match the code contents. See weaknesses section for more details.\n\n* Significance: getting LLMs to adapt to new tasks at a reasonable cost is an exciting endeavour. However, human annotations are still required to emphasize specific chunks of texts. This human-in-the-loop requirement is comparable to that of instruction tuning, but it is not clear whether the results are comparable between the two methods.", "weaknesses": "# Problems in evaluation code\n\nAs can be seen in the code (https://anonymous.4open.science/r/PASTA-10E9/pasta/evaluator.py), the evaluation metrics for JSON Format and pronoun changing tasks do not seem aligned with any intuitive definition:\n\n* the paper describes the metrics for the JSON format task as follows: \"Format accuracy (F. Acc.) measures the accuracy at generating valid JSON. (b) Prediction accuracy (P. Acc.) measures the accuracy at generating the correct target in JSON values after loading the JSON-formatted generations.\". However, in practice:\n  * the evaluation for format accuracy measures whether there is any subsequence of the generated text that is valid json. It is highly debatable whether a model that would generate natural text intermingled with valid json is actually accurate at generating json.\n  * the evaluation for prediction accuracy checks whether the desired string is present anywhere in the json values. As such, an answer that would print out e.g. all of the words in the context into one large json would score well.\n* worse, the evaluation for pronoun change looks only at whether specific pronouns (\"she\", \"he\", \"her\"...) are no longer in the generated text, without verifying that the rest of the text is unchanged. As such, an empty generation \"\" would score 100. This explains partially Figure 3b, where the model's fluency seems to be decreasing steadily as more heads are steered, while the pronoun change accuracy keeps increasing.\n\nThese flaws in evaluation are troubling: a high score on these metrics may be paired with a substantial decrease in overall fluency and perceived quality of the model. This does not seem in line with the paper's stated contribution which is to \"enhance an LLM’s ability to follow user instructions\".\n\n# Overlap with instruction-tuning\n\nThe main focus of the paper is to get models to generate text that follows specific instructions. While the authors cite existing publications on instruction-tuning, their claim that \"PASTA can be used in addition to these approaches to improve some aspects of model steerability\" would need to be substantiated. Indeed, one could assume that instruction fine-tuning will alter attention weights to have the model attend closely to the operator's instructions. The notion that PASTA improvements would transfer to such models would thus need to be tested. Furthermore, the cost of the model profiling phase should be compared to the cost of a lightweight instruction-tuning setup.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)