Pay Attention to What Matters

## Abstract

Abstract Despite the remarkable success of Large Language Models (LLMs), they still exhibit a limited capability to align their outputs to the user instructions. In this work, we introduce a simple and effective method, which we name GUIDE, that mechanistically increases attention scores in instruction tokens. To support this operation, we present Influence , a novel metric that highlights how the user’s instructions propagate through the transformer layers and impact the LLM output. Our results show that GUIDE improves the accuracy of following instructions 29.4 % percent 29.4 29.4\% 29.4 % to 60.4 % percent 60.4 60.4\% 60.4 % , outperforming natural prompting alternatives and Supervised Fine-Tuning up to 1M tokens. 1 1 1 Our experiments are available on https://github.com/netop-team/pay_attention_experiment .

## 1 Introduction

Large Language Models (LLMs) are currently the state-of-the-art of most NLP tasks. Despite this success, pretrained LLMs sometimes struggle to accurately interpret diverse users’ instructions and may generate outputs that do not align with human expectations. Additionally, LLMs may produce biased or hallucinated facts, which can limit their practical usefulness.

Previous work (Kuratov et al., ; Lu et al., ) indicate that transformers are less prone to align with instructions as the context length grows (Kuratov et al. (); Lu et al. ()). In such cases, rather than fulfilling the user’s request, the model generates nonsensical text or repeat segments from the prompt.

A common solution to this problem is Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL). However, these approaches are resource-intensive, time-consuming, and sensitive to the specific data and task. Ideally, a more efficient approach would be one that, once implemented, does not require additional training.

In that sense, due to its low cost and broad accessibility, prompt engineering is widely used to align the outputs of LLMs with user preferences. However, this method does not always produce consistent results and can be very unstable, as demonstrated in (Sclar et al., ).

Figure: (a) GUIDE uses tags (such as <!-> <-!>) to know where to focus. It then enhances the importance of highlighted tokens by biasing the attention scores toward them, as shown by the attention matrices above, where each entry represents the impact of a past token (x-axis) on the ongoing token (y-axis).
Refer to caption: x1.png

In this work, we introduce GUIDE (Guided Understanding with Instruction-Driven Enhancements), a novel and systematic approach that allows users to emphasize critical instructions in their prompts. GUIDE enables users to influence the attention given to specific tokens by simply enclosing important text within tags like <!-> <-!> (as shown on Figure 1(a)). These special tags directs the LLM’s focus, which is done by adding a bias to the attention scores toward the tokens they enclose. Our implementation is open-source and designed for seamless integration. Our experiments demonstrate that GUIDE significantly increases the likelihood of the model following key instructions and retrieving crucial information designated by the user, outperforming natural prompting techniques.

While GUIDE does not require additional training, it does necessitate the careful selection of how much to increase attention weights. In our study, we propose default values for certain tasks, but we also recognize the need to quantify these adjustments. To address this, we introduce a novel metric called Influence. This metric measures the importance of specific tokens in relation to instruction tokens within the text, and we use it to determine reasonable values for the increase in attention weights.

At a high level, Influence is calculated by propagating the contribution of the attention map across each layer, weighted by the norm of the corresponding token embeddings. We demonstrate that this metric correlates with the model’s probability of following specific instructions.

To that end, the main contributions of this work are:

- 1.
The introduction of GUIDE: a mechanistic approach for emphasizing instruction tokens, without need of any further computational resources.
- 2.
The introduction of Influence: a non-gradient metric that quantifies the importance of a given instruction over the text, which we use to adjust and understand GUIDE.
- 3.
Release of PayAttentionPipeline: a HuggingFace-based implementation capable of performing generation with GUIDE and computing Influence (as illustrated in Figure 1)(^2^22The code is available on https://github.com/netop-team/pay_attention_pipeline).

## 2 Related work

Alignment and instruction following. Alignment techniques have the objective to align LLM outputs with human preferences. In general, we refer to being Helpful, Harmless and Honest. Model fine-tuning usually aligns the output of the LLMs with human intents using Reinforcement Learning with Human Feedback (RLHF) (Ouyang et al., ), Reinforcement Learning with AI Feedback (RLAIF) (Lee et al., ) or Direct Preference Optimization (DPO) (Rafailov et al., ).
However, these methods have three significant constraints: they require specialized datasets, often with human annotations, thus reducing efficiency; they involve substantial computational complexity and cost due to the need for additional training; and they demand specialized expertise, as successfully implementing the process can be challenging.
Given these constraints, this type of fine-tuning is typically reserved for general-purpose alignment, ensuring that models are Helpful, Harmless, and Honest (Shen et al., ) and are not suited to address end users’ specific needs. Supervised fine-tuning (SFT) with techniques such as low-rank adapters (LoRA (Hu et al., )) offers a more accessible way to customize a model for individual user requirements. However, these techniques still face the same three limitations, albeit to a lesser extent. Consequently, SFT is typically utilized only for very targeted use cases if used at all.

Utilizing LLMs for automated prompt engineering has demonstrated notable performance.
Black-Box Prompt Optimization (BPO) is a sophisticated framework that automatically refines human-written prompts, often unstructured or ambiguous (Cheng et al., ). Similarly, the PE2 framework (Ye et al., ) enhances prompt performance by refining human-written prompts through a comprehensive search process. Although PE2 avoids additional model training, it increases complexity, latency, and cost, limiting its scalability. Both BPO and PE2 are generally designed for broad enhancements in prompt writing. They are not tailored to meet individual users’ specific intentions or needs.

Due to its low cost and large accessibility, prompt engineering is extensively used to align the output of the LLMs with user preferences. This is clearly demonstrated by popular LLM frameworks like the system in (Yang et al., ), which empowers LM agents to tackle software engineering tasks. This system emphasizes crucial instructions through uppercase text and exclamation marks, like "PLEASE DO NOT DO THAT!" or "THE EDIT COMMAND REQUIRES PROPER INDENTATION." Similarly, the AI Scientist (Lu et al., ), a leading system for automated scientific discovery, uses strong directives such as "ABSOLUTELY DO NOT ADD IT AGAIN!!!" to steer the model’s behavior.
These examples, drawn from highly influential frameworks widely used, underscore the pressing need for end-users to signal what matters most to them in order to guide LLMs toward better alignment with their goals. Currently, users rely on prompt engineering and forceful language to achieve this alignment. However, this approach does not consistently deliver positive results as shown in (Sclar et al., ).

In contrast, our proposed method, GUIDE, offers a reliable and systematic approach that enables users to mechanistically highlight critical instructions within the prompt.

Explainability for Transformers. This is a particularly active area of research with many promising research directions. In the context of this work, we primarily focus on methods that attempt to quantify the significance of one token (or set of tokens) of interest.

Gradient metrics such as Relevance and GradCAM (Chefer et al. (); Chefer et al. (); Selvaraju et al. ()) have demonstrated promising results in Computer Vision, NLP and Text-To-Image tasks. Nonetheless, calculating gradients in large models, particularly those exceeding 7B parameters, demands substantial computational resources.

Activation Patching approaches (Meng et al. (); Zhang and Nanda ()) focuses on perturbing the inputs and checking how impactful will this perturbation be, based on where and how is it applied. They propose two types of perturbation: adding a gaussian noise in token embeddings, or perturbing the phrase in a semantical approach (changing important words in the phrase), and they evaluate the importance of each token by the difference of distributions of logits into these with perturbations.

Abnar and Zuidema () introduced Attention Rollout, a method that uses attention weights to measure the impact of one token on another. This technique propagates attention weights through the layers to determine the importance of each token from any layer to the first layer. However, this method is primarily applicable to encoder-only architectures like BERT. In contrast, generative models typically use decoder-only architectures, making attention rollout less effectives.

In this work, we introduce Influence, a simple and computationally efficient metric specifically designed for decoder-only models.

## 3 GUIDE: ( G uided U nderstanding with I nstruction- D riven E nhancements)

In this section, we present GUIDE, a novel and systematic approach that enables users to highlight critical instructions within the text input provided to an LLM. To understand how GUIDE operates, it is essential first to revisit the core mechanism of self-attention, which drives the functioning of LLMs.

### 3.1 Description of the method

Each token $k$ in the input text is initially represented by an embedding, denoted by $E^{(0)}_{k}$, which undergoes progressive refinement through stacked layers of attention. By the time it reaches the final layer, $L$, this embedding $E^{(L)}_{k}$ is expected to encapsulate all the semantic information required to predict the next token, $k+1$.

The process operates as follows: at each attention layer $\ell$, the embedding of a token $k$ is enhanced with the semantic information of past tokens ($i=1,2,\dots,k-1$) and itself. This enrichment occurs through a residual connection, where the embedding $E^{(\ell)}_{k}$ is updated with the output of the attention layer, which consists of a weighted average of the values $V^{(\ell)}_{i}$ of the past tokens. The vectors $V$, known as values, are derived from a simple linear transformation of the embeddings $E^{(\ell)}_{i}$, for $i\leq k$, and are responsible for carrying the semantic information from the past tokens.

The extent to which previous tokens influence the semantic update of the token $k$ is determined by attention logits, denoted by $w_{k,i}^{(\ell)}$. These logits represent the raw, unnormalized relevance scores between a token of interest $k$, called a key, and each preceding token $i\leq k$, called queries. The logits are then passed through a softmax function, which normalizes them to sum to one. The resulting normalized weights are known as attention scores ($\mathbf{A}^{(\ell)}$) and quantify the degree of influence each past token has on the current token’s semantic representation at a given layer. Denoting $U_{k}^{(\ell+1)}:=\text{Attention}^{(\ell+1)}(\mathbf{E}_{k}^{(\ell)})$
, the operations at layer $\ell$ can be summarized as follows:(^3^33For simplicity, we have excluded the normalization and feedforward layers from this explanation.)

$$ $E^{(\ell+1)}_{k}=E^{(\ell)}_{k}+U_{k}^{(\ell+1)}=E^{(\ell)}_{k}+\sum_{i=1}^{k} \mathbf{A}^{(\ell+1)}_{k,i}V^{(\ell)}_{i},$ (1) $$

The logits, and hence the attention scores, are automatically computed by the model. We argue that the end user should be able to influence the level of attention each token receives by explicitly signaling which instructions or pieces of information are critical. By doing so, the user can effectively guide the model to better align with his/her intention. We propose to achieve this by simply adding a bias, denoted by $\Delta$, to the attention logits of the important tokens, i.e., $\bar{w}_{k,i}^{(\ell)}=w_{k,i}^{(\ell)}+\Delta,$
for all tokens $i$ indicated by the user. While this approach is direct, it proves to be highly effective, as demonstrated in the experimental results section.

### 3.2 Calibrating GUIDE

Using GUIDE, the addition of $\Delta$ directly increases the attention the model pays to the tokens of interest, amplifying their influence on the generated output. However, because attention scores must sum to one, this adjustment reduces the attention given to other tokens. If $\Delta$ is set too high, the model might overly focus on the highlighted tokens, which could disrupt the generation process. Therefore, it is crucial to select an appropriate $\Delta$ that balances these effects.

Our experiments suggest that for the Mistral and Gemma-2 models, a $\Delta$ of 2 works well for emphasizing instructions, while a $\Delta$ of 1 is effective for highlighting specific information within the text. Besides, using $\Delta$ values greater than 5 often led to nonsensical outputs (see Appendices B and F).

Although these default values improve performance, the optimal choice of $\Delta$ depends on various factors, including the model, the nature of the task, etc. The most precise way to determine an appropriate $\Delta$ is through hyperparameter tuning on a validation set.

In this work, we also introduce a heuristic approach for calibrating $\Delta$ with just a couple of forward passes. The idea is to match the influence increase from $\Delta$ to a "natural" level that could be achieved through conventional prompting, such as using uppercase (see Figure 3). This calibration requires a metric that evaluates the influence of the selected tokens and tracks how this impact propagates both vertically across the stacked layers and horizontally across successive tokens.

### 3.3 Influence

Let us denote by $\mathcal{\bar{U}}=(x_{1},\dots,x_{n})$ the overall sequence of tokens associated with the user’s query and with $\mathcal{U}=(x_{i},\dots,x_{j})$ the tokens related to the instruction that the user desires to highlight.

To maintain simplicity and minimize computational cost, we avoid using gradient-based metrics to evaluate the impact of a subset of tokens on the overall sequence (for example, see (Chefer et al., ),(Chefer et al., ), and (Selvaraju et al., )). Instead, a more appropriate option appears to be the Attention Rollout method proposed in (Abnar and Zuidema, ). This metric can be easily computed during the forward pass, aligning well with our needs.

The Attention Rollout approach is based on a natural interpretation of attention scores. It postulates that the influence of a past token $i$ on the update of the current token $k$ is quantified by the attention score $\mathbf{A}^{(\ell)}_{k,i}$. The method addresses the residual connection by assuming that in the updated embedding $E^{(\ell+1)}_{k}$, both the previous embedding $E^{(\ell)}_{k}$ and the update vector $U^{(\ell+1)}_{k}$ contribute equally, each having an impact of $\frac{1}{2}$. The vertical and horizontal flow of the impact $\text{R}_{\mathcal{U}}(E_{k}^{(l)})$ of a given token of interest $\mathcal{U}$ on an embedding $E_{k}^{(l)}$ is hence characterized by the following recurrence:

$$ $\displaystyle\text{R}_{\mathcal{U}}(E^{(\ell)}_{k})=\frac{1}{2}\left[\text{R}_ {\mathcal{U}}(E^{(\ell-1)}_{k})+\text{R}_{\mathcal{U}}(U^{(\ell)}_{k})\right]= \frac{1}{2}\left[\text{R}_{\mathcal{U}}(E^{(\ell-1)}_{k})+\sum_{i=1}^{k} \mathbf{A}^{(\ell)}_{k,i}\cdot\text{R}_{\mathcal{U}}(E^{(\ell-1)}_{i})\right]$ $$

We argue that Attention Rollout inaccurately represents the flow of attention, particularly when handling the residual connection. The norm of the past embedding $E^{(\ell)}_{k}$ is typically about 100 times larger than that of the update vector $U^{(\ell+1)}_{k}$ (see Figure 2(a)). .

As context length increases, one would expect that the importance of a subsequence of tokens would decrease, since the model must process more information. However, by assuming equal contributions from $E^{(\ell)}_{k}$ and $U^{(\ell+1)}_{k}$, Attention Rollout significantly overestimates the importance of past tokens.

This error compounds as the context length increases, leading to an inflated impact estimate that increases with the context and hence negatively correlates with the model’s likelihood of following the token of interest, such as adhering to a specific instruction (Figure 2(b)).

To address this issue, we introduce Influence, a new metric designed to quantify the impact flow of a token or a set of tokens of interest $\mathcal{U}$. This metric corrects Attention Rollout by weighting the contributions according to the norm of the vectors:

$$ $\text{I}_{\mathcal{U}}(E^{(\ell+1)}_{k})=\frac{\left(\|E^{(\ell)}_{k}\|\cdot \text{I}_{\mathcal{U}}(E^{(\ell)}_{k})+\|U_{k}^{(\ell+1)}\|\cdot\text{I}_{ \mathcal{U}}(U_{k}^{(\ell+1)})\right)}{\|E^{(\ell)}_{k}\|+\|U_{k}^{(\ell+1)}\|}$ (2) $$

Figure: (a)
Refer to caption: x3.png

More precisely, Influence $\text{I}_{\mathcal{U}}:\mathbb{R}^{dH}\rightarrow\mathbb{R^{+}}$ (where $d$ is the attention head dimension and $H$ is the number of attention heads), is a transformer interpretability metric designed to quantify how sequences of tokens in the user’s query impact each others and relate to the LLM output. It is designed based on the following principles:

#### Initialization:

We initialize the Influence value as 1 for tokens within the instruction $\mathcal{U}$, and as 0 elsewhere. Let $E_{k}^{0}\in\mathbb{R}^{dH}$ be the embedding of token $x_{k}$. Then, the Influence initialization can be formally defined as:

$$ $\text{I}_{\mathcal{U}}(E^{0}_{k})=\mathbbm{1}_{\{x_{k}\in\mathcal{U}\}}.$ (3) $$

#### Propagation Rules:

Given $m$ embedding vectors $E_{1},\dots,E_{m}\in\mathbb{R}^{dH}$, the joint Influence of the instruction tokens $\text{I}_{\mathcal{U}}:\mathbb{R}^{dH}\times\dots\times\mathbb{R}^{dH}
\rightarrow\mathbb{R}_{+}$ is calculated as the average of each individual Influence, weighted by the norms of each embedding, as follows:

$$ $\begin{split}\text{I}_{\mathcal{U}}(E_{1},E_{2},\dots,E_{m})&=\frac{\sum_{i=1} ^{m}\text{I}_{\mathcal{U}}(E_{i})\|E_{i}\|}{\sum_{i=1}^{m}\|E_{i}\|}.\end{split}$ (4) $$

Additionally, we maintain the invariance of Influence to function composition, i.e.,

$$ $\begin{split}\text{I}_{\mathcal{U}}(f(E))&=\text{I}_{\mathcal{U}}(E).\end{split}$ (5) $$

See Appendix A for detailed derivations.

Using Influence, we can calibrate GUIDE by choosing a $\Delta$ that mimics natural attention enhancement, such as writing in uppercase (Figure 3). This can be easily achieved with two forward passes, one with uppercase text and one without. $\Delta$ is then defined as the difference in log-influence between the two versions. It is important to note that our experiments indicate that an instruction highlighted with GUIDE typically has a greater impact on text generation compared to one highlighted with natural prompting, such as using uppercase, even if their influence scores are similar.

Moreover, in the absence of GUIDE, our experiments demonstrate that Influence correlates positively with the likelihood of a set of tokens impacting the model’s output, such as following an instruction (e.g., summarizing in French) or retrieving specific information (e.g., finding a needle in a haystack). Thus, although not flawless, Influence offers a tool of independent value that can be used to compare and predict the impact of different natural prompting techniques.

Figure: Figure 3: Log of the influence across different layers. This illustrates that with an appropriately chosen $\Delta$, GUIDE can effectively replicate—and even further amplify—semantically intuitive instructions, like using uppercase text.
Refer to caption: x5.png

## 4 Experiments

In this section, we evaluate the benefits of GUIDE and Influence using Mistral-7b Instruct (Jiang et al., ) for following key instructions and retrieving crucial information designated by the user. In Appendix D, we conduct the same experiments with Gemma2-2b Instruct Team et al. () with the same values of $\Delta$, and we obtain very similar results.

### 4.1 Description

#### Summarization in French

To evaluate the capability of GUIDE to support LLMs in producing outputs aligned with the user’s query, we perform experiments related to text translation and summarization.
In these experiments, we have used text from OpenWebText (Gokaslan and Cohen, ), chosen for its variety in context lengths. We have divided the dataset into groups based on context length, containing texts from a 500-token window, such as (0, 500], (500, 1000], and so on. From each group, we randomly selected 20 texts and generated 10 summaries for each text using multinomial sampling (Wiher et al., ).

A needle in a haystack.
To evaluate the impact of our approach on the model’s ability to retain information, we have conducted the Needle in a Haystack. This test involves embedding specific information at a particular location within a text and then asking a question related to that information at the end of the text. Our hypothesis is that by adding extra attention to this text, the model’s outputs would improve, as the final representation should be more closely aligned with the information tokens. We have followed the methodology outlined by (Kamradt, ). Specifically, we have inserted specific information, referred to as the "needle" at variable positions within a given text. After this insertion, we have asked a question to the LLM related to the inserted information (see the complete prompt in Appendix E).

To conduct this experiment, we have sampled 200 texts from the OpenWebText (Gokaslan and Cohen, ) dataset, selecting 50 texts for each context window of size 500, ranging from 0 to 6000 tokens. For each text, the needle was inserted at 10 different quantiles (10%, 20%, …, 100%). We placed the needle immediately after a period (‘.’) to maintain the semantic integrity of the text.

#### JSON Generation

To assess the efficiency of GUIDE in generating outputs in a specified format, we have conducted experiments focused on JSON generation. For our inputs, we have used texts from books written between 1510 and 1699, sourced from the BL Books dataset (Labs, ). We have prompted the model to extract and generate key information about each book in a predetermined JSON format, as detailed in E. We have randomly selected 300 books from the BL Books dataset and divided each text into context length windows of 500 tokens, ranging from 0 to 4000 tokens. These text segments were then incorporated into our template, where the Mistral model was expected to generate a JSON output that precisely followed the specified format.

We have inputted special attention into the tokens of Your response should follow exactly this template and we have then evaluated the Jaccard index between the keys of the generated JSON and the schema.

#### Influence

For each of the experiments mentioned above, we have evaluated the relationship between the Influence metric and the probability of obtaining correct outputs. To achieve this, we have calculated both the ROC AUC score and the correlation between the importance of instruction tokens and the last token in the sequence. Then, we have compared these results through non-gradient metrics, such as Attention Rollout and raw attention scores. Our hypothesis is that Influence has a strong positive correlation with the probability of correct outputs.

Given that the ROC AUC is a classification metric, it was necessary to binarize our scores. In the French summarization experiment, we have done this by assigning a score of $1$ to texts in French and 0 0 to those in other languages. In the "needle in a haystack" experiment, a score of $1$ was given to prompts that successfully identified the needle information, while those that did not were assigned a 0 0. Similarly, for the JSON generation experiment, outputs that adhered to the JSON format were assigned a $1$, and those that did not were assigned a 0 0.

### 4.2 Results

#### Summarization in French

We have conducted experiments with GUIDE, biasing attention scores towards the instruction Summarize in French.
Fig. 4(a) shows the observed probability that the LLM summary is in French when using GUIDE and compares the results achieved with the baseline model, with both uppercase and normal prompts, as well as the performance observed when including ‘Important:’ before the prompt instruction.
Our findings show that GUIDE leads to an improvement from $29.4\%$ to $60.4\%$ with respect to the raw model, and that the best result is achieved with $\Delta=2$. Besides, to confirm that GUIDE does not induce a deterioration of the quality of the generated outputs, we compare the summaries generated in French obtained with the raw prompt and the ones obtained with GUIDE. We observed no noticeable degradation. Further details can be found in Appendix B.

As a baseline, we compare the performance of GUIDE to prompt engineering and Supervised Fine-Tuning (SFT) using LORA (the hyperparameters can be found in Appendix C). Figure 4(a) show that using uppercase or adding ‘Important’ on the instruction does not provides notable improvements, consistently underperforming GUIDE, while Figure 4(b) shows that GUIDE outperforms SFT until 1M training tokens. These results confirms that our method is an effective solution for aligning LLMs to instruction following that does not require additional training.

Figure: (a) Probability of outputting a summary in French.
Refer to caption: x6.png

#### Needle in a haystack

Figure 5 shows the probability of outputting the correct phrase over the context length and the position of the needle, respectively. The Mistral model demonstrates stable performance across varying context lengths and needle positions within this window.

As expected, the addition of $\Delta$ to the needle tokens consistently enhances performance from $87.0\%$ to $92.1\%$, with optimal values of $\Delta$ around $1$. We can also note that, on average, the LLM is more effective at retrieving information when it is located at the beginning or the end of the text. This is in accordance with previous results (Kuratov et al., ; Kamradt, ).

Figure: Figure 5: Heatmap of scores in a needle in haystack test.
Refer to caption: x8.png

#### JSON Generation

We measure the Jaccard index between the keys of the generated JSON and the keys on the schema. We observed that the optimal value for
$\Delta$ is approximately 3, resulting in an average score improvement of 30% compared to the raw model (Figure 6).

We also note that in almost every generation the scores were 0 0 or $100\%$. This indicates that most of the time, the generated output was either a perfect match to the requested schema or not in JSON format at all.

Figure: Figure 6: Jaccard Index vs Context length for the JSON generation experiment.
Refer to caption: x9.png

#### Influence

Table 1 shows the correlation and ROC AUC of each metric to correct output. We note that, attention rollout shows a negative correlation and an AUC below $0.5$ in two out of three experiments. This observation supports our initial hypothesis that attention rollout may not accurately reflect the model focus. Also as expected, we also see that raw attention has a random behavior in two of three setups, with ROC AUC scores around $0.5$.

The stronger positive correlation and ROC AUC between Influence and the likelihood of following instructions supports our hypothesis that our metric better quantifies the attention flow in a Transformer than other existing non-gradient metrics.

**Table 1: AUC and correlation of our metric and the probability of generating a correct output**
|  | Metric | ROC AUC | Correlation |
| --- | --- | --- | --- |
|  | Influence | 0.74 | 0.72 |
| Summarizing in french | Attention rollout | 0.24 | -0.35 |
|  | Raw attention | 0.58 | 0.13 |
|  | Influence | 0.62 | 0.12 |
| A needle in a haystack | Attention rollout | 0.55 | 0.10 |
|  | Raw attention | 0.48 | -0.03 |
|  | Influence | 0.63 | 0.23 |
| JSON generation | Attention rollout | 0.31 | -0.29 |
|  | Raw attention | 0.64 | 0.23 |

## 5 Conclusion

While Transformers represent the state-of-the-art in almost all NLP tasks, they often exhibit unexpected behaviors, particularly hallucination, which becomes more pronounced as context length increases. This work explores the capability of these models to align with specific instructions and introduces GUIDE, a mechanical approach for instruction alignment that does not require further optimization. We demonstrate that GUIDE effectively mitigates hallucination in instruction-following scenarios without significantly compromising output quality.

To evaluate the impact of GUIDE across different context lengths, we introduce Influence, a novel metric for Transformer explainability that quantifies the importance of subsequences of tokens within the context. Additionally, we implement both GUIDE and Influence in a Hugging Face-based pipeline, making them publicly available to the community.

## Appendix A Detailed derivations on influence

Let us denoted by $d$ the transformer head dimension, with $H$ the number of attention heads, and with $s$ the context length. Following the propagation of a transformer layer [Vaswani et al., ], the embedding on layer $\ell$. $\textbf{E}^{(\ell)}$ is computed as follows:

$$ $\textbf{E}^{(\ell)}=\text{Linear}\left(\text{Norm}\left(\textbf{E}^{(\ell-1)}+ \text{Attention}^{(\ell)}(\textbf{E}^{(\ell-1)})\right)\right),$ (6) $$

$$ $\text{Attention}^{(\ell)}(\textbf{E})=\textbf{A}^{(\ell)}\cdot V^{(\ell)}( \textbf{E}),$ (7) $$

where, $\textbf{E}^{(\ell-1)}\in\mathbb{R}^{dH\times s}$ is the embedding on layer $\ell-1$, $\textbf{A}^{(\ell)}$ is the attention matrix on layer $\ell$, $V^{(\ell)}:\mathbb{R}^{dH}\rightarrow\mathbb{R}^{dh}$ is a linear function that maps the token embeddings to the values vector, Norm is a normalization function, and Linear is a conventional multilayer perceptron (MLP) function. Then, we can compute the Influence of token $k$, $\textbf{E}^{(\ell)}_{k}$, as follows:

$$ $\begin{split}\text{I}_{\mathcal{U}}(E^{(\ell)}_{k})&=\text{I}_{\mathcal{U}} \left(\text{MLP}\left(\text{Norm}(\textbf{E}^{(\ell-1)}+\text{Attention}^{( \ell)}(\textbf{E}^{(\ell-1)}))_{k}\right)\right)\\ &=\text{I}_{\mathcal{U}}\left((\textbf{E}^{(\ell-1)}+\text{Attention}^{(\ell)} (\textbf{E}^{(\ell-1)}))_{k}\right)\\ &=\frac{\text{I}_{\mathcal{U}}(E^{(\ell-1)}_{k})\cdot\|{E^{(\ell-1)}_{k}}\|+ \text{I}_{\mathcal{U}}(\text{Attention}^{(\ell)}(\textbf{E}^{(\ell-1)})_{k}) \cdot\|\text{Attention}^{(\ell)}(\textbf{E}^{(\ell-1)})_{k}\|}{\|{E^{(\ell-1)} _{k}}\|+\|\text{Attention}^{(\ell)}(\textbf{E}^{(\ell-1)})_{k}\|}\\ &=\frac{\text{I}_{\mathcal{U}}(E^{(\ell-1)}_{k})\ r^{(\ell-1)}_{k}+\text{I}_{ \mathcal{U}}(\text{Attention}^{(\ell)}(\textbf{E}^{(\ell-1)})_{k})}{1+r^{(\ell -1)}_{k}}\end{split}$ (8) $$

where $r^{(\ell-1)}_{k}:=\frac{\lVert E_{k}^{(\ell-1)}\rVert}{\lVert\text{Attention}^
{(\ell)}(\mathbf{E^{(\ell-1)}})_{k}\rVert}$.

Influence is computed recursively over layers, i.e., when we compute the Influence on layer $\ell$, we have already computed the Influence on layers $1,\dots,\ell-1$. This means that $\text{I}_{\mathcal{U}}(E^{(\ell-1)}_{k})$ is already computed, while we still need to compute $\text{I}_{\mathcal{U}}(\text{Attention}^{(\ell)}(\textbf{E}^{(\ell-1)})_{k})$. Developing equation 7:

$$ $\text{Attention}^{(\ell)}(\textbf{E}^{(\ell-1)})_{k}=\sum_{i=1}^{s}\textbf{A}^ {(\ell-1)}_{k,i}E^{(\ell-1)}_{i},$ $$

$$ $\text{I}_{\mathcal{U}}(\text{Attention}^{(\ell)}(E^{(\ell-1)})_{k})=\frac{\sum _{i=1}^{s}\textbf{A}^{(\ell-1)}_{k,i}\lVert E_{i}^{(\ell-1)}\rVert\ \text{I}_{ \mathcal{U}}(E_{i}^{(\ell-1)})}{\sum_{i=1}^{s}\textbf{A}^{(\ell-1)}_{k,i} \lVert E_{i}^{(\ell-1)}\rVert}.$ $$

Then, if we approximate the norm of the embeddings $\lVert E_{i}^{(\ell-1)}\rVert$ with a constant, we obtain a simplified expression

$$ $\text{I}_{\mathcal{U}}(\text{Attention}^{(\ell)}(E^{(\ell-1)})_{k})=\sum_{i=1} ^{s}\textbf{A}^{(\ell-1)}_{k,i}\text{I}_{\mathcal{U}}(E_{i}^{(\ell-1)}).$ $$

With this approximation, equation 8 becomes

$$ $\text{I}_{\mathcal{U}}(E^{(\ell)}_{k})=\frac{\text{I}_{\mathcal{U}}(E^{(\ell-1 )}_{k})}{1+r^{(\ell-1)}_{k}}r^{(\ell-1)}_{k}+\frac{\sum_{i=1}^{s}\textbf{A}^{( \ell-1)}_{k,i}\ \text{I}_{\mathcal{U}}(E_{i}^{(\ell-1)})}{1+r^{(\ell-1)}_{k}}.$ (9) $$

## Appendix B Evaluation of the quality of outputs

In addition to verifying that the LLM summary is in French in Section 4.2, we have also evaluated the quality of the outputs using BERTScore [Zhang et al., ], calculated in comparison to target summaries generated by a Llama 3 70B model [AI@Meta, ].

To highlight the pertinence of BERTScore, in evaluating the quality of the summaries, we show in Fig. 7 the distribution of the observed BERTScore conditioned to the generated text being in French or not. We observe that the distribution for texts generated in French is shifted to the right compared to those not in French, indicating that BERTScore is a suitable metric for assessing the quality of generated texts.

To measure the impact of GUIDE on the quality of the LLM outputs, we have evaluated the winning rate by comparing the quality of the texts generated with and without GUIDE in terms of BERTScore. Specifically, for each pair of texts generated in French $(t_{i,\Delta},t_{i,\text{raw}})$ by GUIDE and the unmodified (raw) model, we have determined which text had a higher BERTScore.
Table 7 shows that for small enough choice of $\Delta$, the quality of the output is not highly affected, with winning rates of $50.5\%$ for $\Delta=0.5$ and $\Delta=1$ and $49\%$ for $\Delta=2$. These results indicate that GUIDE maintains the model’s capability to generate semantically correct text. However, as mentioned in Sec. 3.2, larger values of $\Delta$, e.g. $\Delta=5$ results in poor outputs (see also Appendix F).

Figure: Figure 7: Distribution of Bert scores conditioned to the generated text being in French or not.
Refer to caption: x10.png

## Appendix C Supervised Finetuning hyperparameters

In our supervised fine-tuning experiments, we leveraged LoRa techniques Hu et al. [], setting the sequence length to 8192 using sample packing and block-attention to prevent cross-sample contamination. We configured the LoRa rank to 64 and set the alpha parameter to 16. For regularization, we applied a dropout rate of 0.05. To maximize the adapter’s expressiveness, our LoRa implementation targeted all modules. The batch size was set to 128k tokens, with a maximum learning rate of 1e-4, following a cosine scheduler with a 10-step warm-up. Training was conducted for 2 epochs, tracking the number of tokens processed at each step.

## Appendix D Experimental results with other models

We have conducted the same studies presented in Section 4 also using the Gemma 2 - 2B Instruct model [Team et al., ]. Our results indicate that, even with smaller models, GUIDE can still improve the accuracy of following instructions, increasing the accuracy from $43.4\%$ to $59.8\%$ on summarization in French , $65.2\%$ to $77.5\%$ on retrieval and $14.4\%$ to $24.1\%$ on JSON generation (see Tables LABEL:tab:_summarization-gemma, LABEL:tab:_needle-gemma and LABEL:tab:_json-gemma).

**Table 3: Summarization in French results for Gemma 2- 2B**
| $\Delta$ | Score |
| --- | --- |
| 0.0 | 0.434 |
| Uppercase | 0.408 |
| 0.5 | 0.514 |
| 1.0 | 0.571 |
| 2.0 | 0.598 |
| 5.0 | 0.002 |

## Appendix E Prompts used in experiments

## Appendix F Examples of poor generation of text

This Appendix provides examples of the model failing to follow user instructions.

### F.1 Raw generation (without GUIDE)

A common error occurs when the model simply repeats the given prompt. This example has been obtained using multinomial sampling at generation time and without GUIDE ($\Delta=0$).

The LLM sometimes generates a very limited number of tokens that, although present in the prompt, are nonsensical when presented alone (and do not fulfil the user task). In this instance, we employ multinomial sampling to generate the text (without GUIDE).

### F.2 Generation with GUIDE using unreasonably large Δ Δ \Delta roman_Δ

We have noted that $\Delta>5$ very often produces incorrect/nonsensical output. In the following example, we have set $\Delta=9$ and performed greedy decoding.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Here are some questions that I think will clarify the purpose of the proposed method:

1) Does SFT achieve similar effect on the attention distribution (I suspect not) than directly augmenting them? If not, you can further demonstrate the superiority of the method as it may generalize on new instruction much better than SFT.
2) Details of needle in a haystack experiment is a bit ambiguous. Is the "attention token" added across the needle token or the instruction tokens? If it's the needle token, it does not make much sense to demonstrate this as an augmentation to the instruction-following skills as you are essentially marking out the key information to the answer here.
3) What does the y-axis label (Depth) of Fig 4 in section 4 mean? Could you provide a clearer definition in the section about what is 'Depth' and how is that related to the task of 'needle in a haystack'?

Here are some questions for the future exploration:

1) Solution to lost in the middle: with something like a normal distribution of delta, could you use attention augmentation as a general solution to "lost in the middle" for long context inference? 
2) For LLMs trained with structured prompt that highlights the instruction tokens in the input (Llama3-instruct for instance), would the method still work? 

I am happy to adjust my evaluation if the questions and weaknesses are discussed. I think it is an interesting paper that just needs to be polished a bit more.
soundness: 2
strengths: 1) The presentation of the methods are fairly straightforward and clean.
2) The proposed method is very intuitive and simple.
3) Without any extra resources, a substantial improvement has been seen over all the selected datasets, making it potentially an appealing plug-and-play mechanism to be added to deployed model to highlight crucial part of the input.
summary: This paper presents 1) an inference-time attention enhancement mechanism to improve LLM's instruction-following capability given user's annotation (which part of the input is crucial instruction) 2) *Influence*: a simple yet reasonable improvement to *Attention Rollout* that considers the ratio between the residual stream and the attention output. 

The proposed pipeline is constructed using 1) and 2). Specifically, for each input with marked instruction, the *influence* scores when the marked instruction tokens is capitalized is calculated, and a delta weight is added to the QK attention weights before SoftMax at the marked instruction tokens to mimic the same *influence* score to boost the importance of the tokens in the context.

The efficacy of the proposed pipeline is evaluated over three settings: French summarization, needle in a Haystack, and JSON generation. All results indicate that a well-calibrated attention weight augmentation can improvement performance of the LLM on the target task, outperforming the prompting baselines.
weaknesses: Sec 3:
1) The concept of 'influence' as an enhancement to attention rollout is intriguing and initially seems valuable. However, the exact role of influence within the implementation remains somewhat ambiguous. At first the glance, the narrative suggests that influence is computed on a per-example basis to adjust the attention dynamically (delta). Yet, the inclusion of plots and tables depicting a consistent delta value across examples implies that this parameter could potentially be determined via straightforward heuristics or a hyper-parameter search. This raises questions about the essentiality of influence in the proposed method. If influence is not critical to the functionality, a direct comparison of influence and traditional attention rollout—perhaps in basic settings like verb singularity classification—would be imperative to justify its mention as a key contribution. In summary, my recommended experiments would be:
- Add another line for the experiments with "flexible delta based on influence"
- Comparing the performance of influence and attention rollout on the task of localizing the attributed attention for singular verb classification with the same setting than the original attention rollout paper.

Sec 4:
1) The experiments seem to employ a limited set of prompt templates, which might not fully capture the model's versatility in handling varied instructions. The effect of different prompt structures and the positioning of instructions within sentences could potentially alter the outcomes, suggesting a need for a broader range of templates in evaluation.

2) The choice of p(output in French) as the only metric, while straightforward, does not address the quality of the summarization. Considering the discussion on potential issues with high attention augmentation (e.g., delta=2 leading to nonsensical outputs), it is vital to also assess the quality of the generated summaries, either through human evaluations or automated metrics, to ensure the practical utility of the model outputs.
 
3) holistic evaluation of the impact: similar to weakness 4.2, there is a need for a more comprehensive evaluation of how attention augmentation influences the model’s overall comprehension and performance. Specifically, how does attention augmentation affect the comprehension of the other part of the input? Does it degrade the general performance? Could it be used my malicious users to hack and jailbreak the aligned LLM by forcing attention at unintended tokens? I understand that these evaluations may require a lot of effort but I believe they are crucial questions to be answered to claim the method to be truly useful. Here are some proposed experiments that you may try:
- comprehension: Add a dataset that requires complex reasoning to see if paying heavy attention to the instruction would break the performance.
- safety: Pick an aligned model and test it on an adversarial attack benchmark. Two ways of attack by hacking the attention augmentation: 1. Wrapping the adversarial instruction with the attention special token. 2. Wrapping a random 2/3 of the tokens in the input to significant increases the temperature.

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 3
questions: Please see "weakness". 

Minor comment : the citation format seems to be wrong for in-place citation. You should use `\citet` for those and it should look like "... appears to be the attention rollout method proposed in Abnar and Zuidema (2020)."
soundness: 2
strengths: (1) The proposed method is very simple and intuitive. The experiments clearly showed that the methods improved "instruction following" at a coarse granularity. Instruction following is an important topic in modern language model development and the proposed method could have a significant impact and adoption in real applications.

(2) The proposed influence score is an intuitive way to guide the setting of the attention bias.
summary: This paper tackles the instruction following problem and proposes an attention manipulation method to force the model to pay attention to the user instructions or important information in the context. The method is very simple: simply add a bias to the attention logits (before normalization)  for the instructions/important information. The authors also developed an "influence" score to guide the setting of the bias. Inference is a score that describes how much one part of the text affects the later part of the text. The authors set the bias such that the influence matches that when setting the instructions to be upper case.

The experiments are mostly synthetic or toy to clearly demonstrate the instruction following ability: summarization in French (metric: p(output in French), needle in a haystack, and JSON generation (Jaccard index of the keys). The experiments demonstrated the effectiveness of the proposed method, improving over both baselines and uppercasing the instructions.
weaknesses: (1) Since this is a training-free method, it would be even better if the authors conducted the experiments on more models with different scales. Given that the experimented models are not very big (<7B) and strong (there are better models now like Llama-3), the lack of comprehensive experiments with different models may undermine the credibility of the experiments.

(2) The tasks the authors selected are rather synthetic. Moreover, the metrics are too simple: for example, the correctness of the French translation and the information extracted for JSON generation should also be evaluated. Even though the focus of the method is instruction following, the correctness of the output should also be measured to understand if the method hurts the performance. Also, the authors should make the evaluation part clearer: for example, I did not find how "whether output is French" is evaluated. 

(3) There are several existing instruction-following benchmarks, such as IFEval (Zhou et al.). Even though the existing instruction following benchmarks are also quite synthetic (mostly focusing on writing tasks), the instructions are at least more complex and the evaluation is more rigorous. 

Zhou et al. Instruction-Following Evaluation for Large Language Models

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: - Could you please describe your work with respect to Zhang et al. (ICLR 2024) as mentioned in the weakness?
- Could you measure the change in output quality as you apply your method?
soundness: 2
strengths: - The paper is tackling an important problem in LLM. For instance, as the user prompt gets longer, we observe the LLM not following every part of the user instruction. A method that can enhance instruction following is a great contribution to the field.
- The paper is clear and generally well written.
- Influence score seems to give a good insight into evaluating the impact of tokens on the entire sequence.
summary: The paper introduces GUIDE, a method for improving Large Language Models' ability to follow instructions by emphasizing important tokens in prompts, using special tags that bias attention scores. GUIDE requires no additional training, making it efficient and accessible. To support this, the authors also introduce the Influence metric, which quantifies the impact of instructions on the model's output. Experiments show GUIDE significantly enhances alignment with instructions across tasks, outperforming other methods like fine-tuning​.
weaknesses: - It is missing a very relevant work from ICLR 2024, "Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs" by Zhang et al.
- While it is good to know that the llm is now producing French more often, but this may be at the cost of degraded output accuracy. I suspect that if the attention weights are manually changed, the quality of the outputs may be affected. I think it is very important to measure this as well.
- I think Influence metric shouldn't be called "metric", because a low influence score doesn't necessarily mean the model is not following the instruction well. I think "probe" is a more appropriate term, which helps us to perform intrinsic analysis.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 1
questions: - Is it possible to propose an algorithm to automatically improve the methodology using the proposed metrics? The intent of the question is related with a suggestion to improve the contribution of this study.
- As the title “PAY ATTENTION TO WHAT MATTERS” suggests, it is important to put the attention on the important parts, but in some cases, the instruction sentence is not the only important part. For example, in the case of sentence summarization, it is necessary to give strong attention to important parts of long sentences instead of instruction sentences. What is the relationship between that case and this study?
- I assume that the experiment was done with zero-shot prompting (but it was unclear because of the lack of description), what are the differences if we conduct experiments with few-shot in-context learning?
soundness: 2
strengths: - This study introduces a prompt engineering method for large language models which increases attention scores in instruction tokens by simply enclosing important text within tags like <!-> <-!>.
- A new metric is proposed to clarify that the proposed method focuses on the instruction text that are highlighted by the proposed method.
summary: - This study introduces a prompt engineering method for large language models which increases attention scores in instruction tokens by simply enclosing important text within tags like <!-> <-!>.
- In addition, the study presents a metric that highlights how the user’s instructions propagate through the transformer layers and impact the LLM output.
- Experimental results show that the method improves the accuracy of following instructions.
weaknesses: - There is already a prior research on prompt engineering method that shows that enclosing the tags improves performance, but this study does not describe the relationship. e.g. https://arxiv.org/abs/2309.13078
- Assessing the impact of attentions from the lens of vector norms is already a hot topic, but there is no description about the relationships.
https://arxiv.org/pdf/2004.10102
- Lack of relationships between sections, mathematical symbols, and descriptions of key points in each section makes the manuscript difficult to read. For example, In 3.1 DESCRIPTION OF THE METHOD, which is a preliminary setup and which is the proposed method? Another example is that in 3.3 INFLUENCE, R_U function and I_U function at the first layer are not defined so that we cannot understand what those functions mean when we are given only the recursive formula.