Differential Transformer

Questions part:

1. Can the authors elaborate further on why we need a fixed multiplier $(1-\lambda_{\text{init}})$ in DIFF Transformer Equation (3)?

The related part in Equation (3) is:

The multiplier $(1 - \lambda_{\text{init}})$ can be considered as the initialization of the scale parameter of $\operatorname{LN}$, designed to align the gradients with Transformer. At the early training stage, with differential attention defined in Equation (1), for a token $x$ in $\mathrm{head}_i$ we have:

Consider the gradient of $\operatorname{LN}$:

With the multiplier $(1 - \lambda_{\text{init}})$, we counteract the gradient factor $\frac{1}{1-\lambda_{\text{init}}} > 1$, better aligning gradients of Diff Transformer with Transformer to avoid hyperparameter adjustments.

The detailed gradient derivation can be found in Appendix E.

2. In section 3.7, Is absmax quantization the suitable approach to adopt in this context? Why not consider quantization approaches that are robust to activation outliers, such as LLM.int8()?

Recent research suggests that activation outliers are closed related to attention mechanism both for attention logits [1] and hidden states [2]. In section 3.7, we aim to investigate whether activation outliers can be reduced natively through attention design. To illustrate the benefits of reducing outliers, we primarily employ absmax quantization. Advanced quantization techniques can also be applied to DIFF Transformer, and the improvement over Transfomer still exists as activation outliers are reduced. We are also exploring quantization-aware training with DIFF Transformer, which shows promising preliminary results.

[1] Quantizable transformers: Removing outliers by helping attention heads do nothing. Bondarenko et al. 2023.

[2] Massive activations in large language models. Sun et al. 2024.

We appreciate the reviewer finds our approach promising and experiments solid. Here are our responses:

Weaknesses part:

1. Related work can be discussed more thoroughly in the paper. For example, [1] discusses that Transformer models tend to mis-allocate attention to irrelevant contexts potentially biased by their position within the context, and [2] also shows LLMs can be easily distracted by irrelevant contexts.

We appreciate the compelling related works that the reviewer mentions. In [1], the authors identify that Transformer exhibits an U-shaped attention bias. It tends to misallocate attention to irrelevant contexts at the beginning and end of the sequence. Similarly, [2] investigates the distractibility of Transformer models and finds that performance significantly drops when irrelevant information is introduced into the context. Our observations align with the conclusions of both studies, confirming that Transformer is susceptible to distraction from irrelevant information, and can misallocate attention scores to these terms.

We will add this to discussion section when there are extra pages available.

2. While the proposed DIFF Transformer shows promising performances, they have to be trained from scratch which can be computationally expensive compared to post-training approaches that mitigate attention noises in standard Transformer models, such as the calibration technique used in [1], and improved decoding approaches used in [3] and [4]. Currently, vanilla Transformer is the only baseline considered in the paper. However, I would appreciate if the authors can show how DIFF Transformer compares to other more lightweight techniques that fix Transformer's attention noises.

[1] mitigates the U-shaped attention bias through a calibration mechanism. [3] utilizes a contrastive method in output distribution, named as context-aware decoding, to improve the faithfulness of LLM. [4] proposes to use the ratio of attention weights as a hallucination detection model and guide decoding with it. These are all effective works that identify problems in LLM similar to attention noise, and mitigate it with lightweight techniques.

In DIFF Transformer, we focus on solving the attention noise problem fundamentally. Training models from scratch helps to fully unleash the potential of the proposed approach. By training from scratch, we can comprehensively enhance DIFF Transformer's capabilities so as to thoroughly validate the benefits of canceling attention noise. Moreover, it is also feasible to apply DIFF approach to pretrained LLMs with post-training methods. It is an interesting and promising topic to explore in the future.

We will add this to discussion section when there are extra pages available.

[1] Found in the Middle: Calibrating Positional Attention Bias Improves Long Context Utilization. Hsieh et al. 2024.

[2] Large Language Models Can Be Easily Distracted by Irrelevant Context. Shi et al. 2023.

[3] Trusting Your Evidence: Hallucinate Less with Context-aware Decoding. Shi et al. 2023.

[4] Lookback Lens: Detecting and Mitigating Contextual Hallucinations in Large Language Models Using Only Attention Maps. Chuang et al. 2024.

3. Apart from pretraining DIFF Transformer, there is no further instruction tuning experiments included in the current paper. I would be interested in seeing whether instruction tuned DIFF Transformer also shows better performances than standard instruction tuned Transformers.

We believe DIFF Transformer is also effective in instruction tuning experiments. Take hallucination evaluation as an example: in addition to question answering and text summarization, multi-turn dialogue is also an important scenario where dialogue history often districts generation of the next response in instruction-tuned LLMs. By mitigating contextual hallucination, DIFF Transformer also has the potential to perform well in this task.

We are interested in exploring this topic further in the future. We will include updates in the revised version.

Thank you to the authors for the response. I hope the authors would include the discussions here into the revision, and I look forward to follow-up results on applying DIFF in a post-pretraining/instruction tuning paradigm.

We are glad the reviewer finds our approach innovative and promising. We also appreciate the comment regarding computational efficiency. We list some possible ways to further improve effiency:

• Impact of Model Size: As the model size increases, the computation occupied by FFN grows. The gap of throughput of DIFF Transformer will be reduced in larger models, which can also be observed in Table 7.

• Advanced Kernels: Efficiency of DIFF Transformer can be further improved by implementing with most advanced FlashAttention techniques such as FlashAttention3 [1] or with kernel implementation designed for differential attention.

[1] Flashattention-3: Fast and accurate attention with asynchrony and low-precision. Shah et al. 2024.

We appreciate that the reviewer finds our paper solid and insightful. Here are our responses.

Weaknesses Part:

1. The authors mention "promoting the emergence of sparse attention patterns" multiple times, but do not provide statistics and quantification of the sparsity of attention distribution between DIFF Transformer and the general Transformer.

We provide statistics of attention scores to demonstrate the sparse attention pattern of DIFF Transformer in Table 3. The statistics are collected from the key information retrieval task. "Attention to Answer" denotes the sum of attention scores that the generated tokens allocate to the answer span. "Attention Noise" denotes the sum of scores that the generated tokens allocate to the irrelevant noise context.

Compared with Transformer, the generated tokens of DIFF Transformer allocate much higher scores to the answer span, and lower scores to the noise context. Further given that the answer span only includes a few tokens, and the attention scores on the noise context (about 4K tokens) are almost uniformly distributed, the statistics can prove that DIFF Transformer attends to key information more sparsely.

It would be better to clearly point out the relationship of this table to the emergence of sparse attention pattern. We will add this in the revised version.

2. The authors did not discuss or conduct experimental comparisons with work related to the sparse attention.

4. The absence of a Related Work section limits the paper's ability to contextualize its contributions within the broader field.

We appreciate the reviewer’s suggestions for the related work. We compare and discuss them in the response to the fifth question.

3. It would be even better if the effectiveness of the DIFF Transformer could be validated on image or speech modalities.

We are conducting experiments on more modalities with DIFF Transformer these days, and there are some positive results. We will add results of more modalities in the revised version.

Questions Part:

1. Do both attention calculations in differential attention use RoPE positional encoding?

Yes.

2. Similar to multi-query attention, can $K_1$, $K_2$, or $Q_1$, $Q_2$ share parameters?

Yes. We conduct experiments on models with 300M parameters. The validation loss of Transformer is 4.406. For a full-parameter DIFF Transformer, the loss is 4.343. For a DIFF Transformer with $Q_1=Q_2$, the loss is 4.347. For a DIFF Transformer with $K_1=K_2$, the loss is 4.346. The performance indicates it is possible to share $K_1$, $K_2$, or $Q_1$, $Q_2$ in DIFF Transformer.

3. Is there a significant difference in sparsity between DiffAttn and regular attention? Does the first term of DiffAttn resemble the pattern of standard attention (implying that the second negative attention term serves to cancel out noise)?

Yes. We demonstrate the sparse pattern in Figure 1 and Table 3. We further explain it in the response to the first question in the Weaknesses Part above.

We visualized two attention maps. The sum of the second attention term equals to $\lambda$. For most layers, $\lambda < 1$. Under this circumstance, we find on noise context, the second term allocates almost same scores as the first term, while on key information, the second term allocates less scores than the first term. Therefore after subtraction, the noise is canceled out, and the key information is kept. For a few layers with $\lambda > 1$, the role of the two terms exchanges, and the conclusion is the same.

4. What is the calculation formula for Table 3?

Firstly, we define the token set that contains all key information tokens as $A$, set contains all noise context tokens as $N$, and set contains the generated tokens as $G$.

For conventional attention, we define attention score, i.e. the output of $\operatorname{softmax}$ function, as $a_{ij}$, which means the score that the $i$-th token allocates to the $j$-th token, $1 \le j \le i$.

"Attention to Answer" is calculated with:

"Attention Noise" is calculated with:

For differential attention, we define attention scores of the first term as $a_{ij}^{(1)}$ and scores of the second term as $a_{ij}^{(2)}$.

"Attention to Answer" is calculated with:

"Attention Noise" is calculated with:

In the DIFF Transformer model we report, $\lambda \in (0,1)$ stands for all layers, so $1-\lambda$ is always positive.

The scores are then averaged across all layers and heads to obtain the final scores in Table 3. The calculation formula ensures all reported scores are within range $[0,1]$.

5. Please discuss the relationship between differential attention and sparse attention.

The mentioned works are both influential paper. Routing Transformer [1] assigns queries and keys to clusters and performs attention only within each cluster. The complexity of attention is reduced to $\mathcal{O}(n^{1.5})$ with the proposed method. Sparse Transformer [2] uses sparse factorizations of the attention matrix to reduce the complexity of attention to $\mathcal{O}(n^{1.5})$. Both works create sparsity by restricting attention to a subset of tokens to accelerate attention. The performance of such methods is usually worse than their full attention counterparts due to the discarded tokens. It is also difficult for them to perform well in tasks such as key information retrieval.

In comparison, DIFF Transformer aims to natively create sparsity in full attention. DIFF Transformer can attend to important context accurately and sparsely without discarding any tokens, enabling it to outperm full attention Transformer in language modeling and tasks such as key information retrieval.

We will add this to discussion section when there are extra pages available.

[1] Efficient Content-Based Sparse Attention with Routing Transformers. Roy et al. 2021.

[2] Generating Long Sequences with Sparse Transformers. Child et al. 2019.

Thank you for your detailed response. I have no further questions. I will keep my very positive score.

Moreover, I would like to suggest modifying "promoting the emergence of sparse attention patterns" to "promoting the emergence of concentrated attention patterns". The term "sparse" implies that some attention scores are zero (which benefits sparse matrix operations), the "concentrated attention" seems more appropriate to describe the effect of differential attention.

We are glad the reviewer finds our idea well-motivated and our experimental results strong. We respond to the reviewer’s questions below.

1. Why the learnable $\lambda$ is re-parameterized in Equation (2)?

$\lambda$ is multiplied to $\operatorname{softmax}$, where $\operatorname{softmax}(\mathbf{q}, \mathbf{k}) = \frac{\exp(\mathbf{q} \cdot \mathbf{k})}{\sum \exp(\mathbf{q} \cdot \mathbf{k})}$. Parameters in $\lambda$ learn with the same learning rate as other parameters in the model, therefore $\lambda$ should take a similar formulation as $\operatorname{softmax}$ to synchronize the learning dynamics.

Consider the re-parameterization $\lambda = \exp(\mathbf{\lambda_{q}} \cdot \mathbf{\lambda_{k}}),~ \mathbf{\lambda_{q}}, \mathbf{\lambda_{k}} \in \mathbb{R}^d$, we then have $\lambda \cdot \operatorname{softmax}(\mathbf{q}, \mathbf{k}) = \frac{\exp(\mathbf{q} \cdot \mathbf{k} + \mathbf{\lambda_{q}} \cdot \mathbf{\lambda_{k}})}{\sum \exp(\mathbf{q} \cdot \mathbf{k})}$ , which synchronizes the learning dynamics of parameters in $\lambda$ with model parameters in $Q, K$ projection. Moreover, to assign $\lambda$ an initialization value $\lambda_{\text{init}}$ and enables it to learn values smaller or larger than $\lambda_{\text{init}}$, we give the final formulation in Equation (2):

2. Why in Equation (3) there is a term $1 - \lambda_{\text{init}}$ for each head?

The related part in Equation (3) is:

The multiplier $(1 - \lambda_{\text{init}})$ can be considered as the initialization of the scale parameter of $\operatorname{LN}$, designed to align the gradients with Transformer. At the early training stage, with differential attention defined in Equation (1), for a token $x$ in $\mathrm{head}_i$ we have:

Consider the gradient of $\operatorname{LN}$:

With the multiplier $(1 - \lambda_{\text{init}})$, we counteract the gradient factor $\frac{1}{1-\lambda_{\text{init}}} > 1$, better aligning gradients of Diff Transformer with Transformer to avoid hyperparameter adjustments.

The detailed gradient derivation can be found in Appendix E.

3. Why the RMSNorm for each head is so important for DIFF Transformer stability?

The importance of headwise RMSNorm can be explained from two aspects:

• The sum of attention scores is $(1 - \lambda)$ in DIFF Transformer, while it is $1$ in Transformer. The RMSNorm normalizes the output value of differential attention and ensures each token in each head has moderate magnitude.

• We find that the headwise RMSNorm plays an important role in sparse attention. As differential attention exhibits sparse attention pattern, the variance in output across different tokens and different heads can be larger than in the conventional attention. Under this circumstance, we find normalizing different tokens in different heads to the same scale is beneficial.

Hi Minh Nguyen Hoang,

In reviewing previous cosine-based attention approaches, Cosformer [1] introduces a cos-based re-weighting mechanism as a relative positional encoding to enforce locality in linear attention. SwinV2 [2] applies the cosine operation to attention logits before softmax to stabilize training.

The goal of the cosine operation is to promote locality or stabilize training. Directly replacing softmax with a cosine operation in attention usually leads to worse performance compared to conventional attention.

[1] cosformer: Rethinking softmax in attention. Qin, et al. 2022.

[2] Swin transformer v2: Scaling up capacity and resolution. Liu, et al. 2022.

Hi Xinyu Zhao,

In the paper, the attention weights come from all answer tokens. In the key information retrieval task, we observe that all answer tokens are answer-relevant tokens. We find averaging the weights across all answer tokens results in more stable and accurate outcomes. Calculating with the first answer token can yield similar conclusions as in the paper.