DINT Transformer

Paper: arXiv
Authors: Yueyang Cang, Yuhang Liu, Xiaoteng Zhang, Erlu Zhao, Shi Li
Date of Publication: 29th January 2025

Overview

This paper introduces a new transformer architecture based on differential transformer. Authors claim that a major drawback of differential transformer is that it’s rows of attention matrix aren’t normalised (they don’t sum up to one) so it causes numerical stability. And also authors claim that differential transformer lack global context modeling (they don’t explain why or how, they just declare it). To solve these issues they come up with a new architecture involving an ‘integration’ mechanism that apparently fixes this.

Differential Attention

The original differential transformer is a very good paper (must read, I’ll maybe make a post and link it here in future). I’ll briefly recap the architecture here. The authors of that paper introduce the differential attention mechanism. Authors calculate 2 attention matrices instead of one, and subtract one attention activation matrix from another (they take the difference, hence the name differential transformer). The authors came up with this as a way to reduce noise in the attention matrix. Intention was that, by allowing the model directly to subtract an attention matrix from another, it can learn to subtract the unnecessary noise. It was pretty popular at the time.

Specifically, given $X\in {\mathbb{R}}^{N\times d_{model}}$, it is projected to $Q$,$K$,$V$ matrices as usual. But,

$$[Q_1;Q_2]=X{W}_Q,[K_1;K_2]=X{W}_K,V=X{W}_V$$

$\text{where, }{Q}_1,Q_2,K_1,K_2\in {\mathbb{R}}^{N\times d}\text{ and }V\in {\mathbb{R}}^{N\times 2d}\text{ and }{W}_Q,W_K,W_V\in {\mathbb{R}}^{d_{model}\times 2d}$

We get 2 query matrices, 2 key matrices. We calculate 2 different attention matrices.

$$\text{DiffAttn}(X)=\left(\text{softmax}\left(\frac{Q_1{K}_1^{\top }}{\sqrt{d}}\right)-\lambda \cdot \text{softmax}\left(\frac{Q_2{K}_2^{\top }}{\sqrt{d}}\right)\right)V$$

$\text{where, }$

$$\lambda =\exp ({\lambda }_{q1}\cdot {\lambda }_{k1})-\exp ({\lambda }_{q2}\cdot {\lambda }_{k2})+{\lambda }_{init}$$

${\lambda }_{q1},{\lambda }_{k1},{\lambda }_{q2},{\lambda }_{k2}\in {\mathbb{R}}^d$ are learnable vectors and ${\lambda }_{init}\in (0,1)$ is a constant used for initialization.

Differential Transformer is a good paper.

Group Norm is done along with layer norm too (I’ll go a bit in depth later while explaining DINT). A fixed multiplier $(1-{\lambda }_{init})$ is used after GroupNorm, which aligns the gradient flow with Transformer. And then there’s a linear layer after multi-head Differential Attention to project back to $d_{model}$

DINT Attention

We had differential attention with no differential equations. Now we have differential integral attention with neither differential equations nor integration.

DintAttn Algorithm

In DINT, we calculate the two attention matrices of Differential Attention as usual. Let’s call them $A_1$ and $A_2$

$$A_1=\text{softmax}\left(\frac{Q_1{K}_1^{\top }}{\sqrt{d}}\right)$$ $$A_2=\text{softmax}\left(\frac{Q_1{K}_1^{\top }}{\sqrt{d}}\right)$$ $$A_{diff}=A_1-\lambda A_2$$

The ‘integral’ component of DINT computes the average attention scores of $A_1$‘s columns.
Essentially $G\in {\mathbb{R}}^{1\times N}$, where each element is the average value of each of the $N$ columns of $A_1$

$$G_{\text{expanded}}=\left[\begin{array}{c}G\\ G\\ ⋮\\ G\end{array}\right]\text{where },G_{expanded}\in {\mathbb{R}}^{N\times N}$$

We just repeat the row vector $G$, $N$ times and stack to create $G_{expanded}$.

$$\text{DINTAttn}(X)=(A_{diff}+\gamma \cdot G_{expanded})V$$

We set $\gamma$ as $\lambda$. And if $\gamma =\lambda$, the rows are normalised and sum up to 1.
I’m being honest, I don’t get why this works, but it just works, and the rows are normalised and sum upto 1. I have checked it numerically.

The authors claim that $G$ captures the global important features.

Multi-head DINT Attention

$$\begin{array}{c}{\text{head}}_i=\text{DiffAttn}(X;W_Q^i,W_K^i,W_V^i,\lambda )\\ \\ {\overline{\text{head}}}_i=\text{LN}({\text{head}}_i)\\ \\ \text{MultiHead}(X)=\text{Concat}({\overline{\text{head}}}_1,\dots ,{\overline{\text{head}}}_h)W_O\end{array}$$

where, $\lambda$ is shared between all the heads (that is the case in differential attention too).
$W_O\in {\mathbb{R}}^{d_{\text{model}}\times d_{\text{model}}}$ is a learnable projection matrix.
$LN$ uses RMS Norm for each head.

Even after headwise normalization, there is a group norm for more stable training.

def DintAttn(X, W_q, W_k, W_v, λ):

1. Q1, Q2 = split(X x W_q) // Split matrix multiplication result for Queries
2. K1, K2 = split(X x W_k) // Split matrix multiplication result for Keys (Corrected to W_k)
3. V = X x W_v // Calculate Value matrix
4. s = 1 / sqrt(d) // Scaling factor
5. A1 = softmax((Q1 x K1.transpose(-1, -2)) x s) // Attention weights 1
6. A2 = softmax((Q2 x K2.transpose(-1, -2)) x s) // Attention weights 2
7. A3 = repeat(average(A1, column), n) // Average Attention (n is the number of rows in X)
8. return (λ * A3 + A1 - λ * A2) x V // Final Attention Output

def MultiHeadDINT(X, W_q, W_k, W_v, W_o, λ):

1. Initialize an empty list O // For outputs of each head

2. for i = 1 to h: // Iterate through heads

   • Get the i-th slice of weight matrices: W_qi, W_ki, W_vi // Slicing for the i-th head
   • O_i = GroupNorm(DintAttn(X, W_qi, W_ki, W_vi, λ)) // Differential Attention for i-th head
   • Append O_i to list O

3. return Concat(O) x W_o // Concatenate head outputs and multiply by output weight

DINT Transformer Layer

Pretty standard transformer decoder layer. We use RMS norm instead of Layer Norm. SwiGLU is used as activation function and of course there’s DINT attention.

$$\begin{array}{rl}{Y}^l& =\text{MultiHeadDINT}(\text{LN}(X^l))+X^l\\ X^{l+1}& =\text{SwiGLU}(\text{LN}(Y^l))+Y^l\end{array}$$

Results

As per authors this new architecture beats Differential Transformer in every single benchmark. They test on context based tasks like needle in a haystack, in-context learning, summarization, question answering. They evaluate the scalability, they compare loss of DINT with Differential Transformer as parameters and tokens are scaled. They analyse attention matrices.

Pretty much in every single metric it beats differential transformer. Read the paper for specific details.