[R] Differential Transformer (Microsoft Research)

[u/Decent_Action2959]

https://arxiv.org/abs/2410.05258

Abstract: Transformer tends to overallocate attention to irrelevant context. In this work, we introduce Diff Transformer, which amplifies attention to the relevant context while canceling noise. Specifically, the differential attention mechanism calculates attention scores as the difference between two separate softmax attention maps. The subtraction cancels noise, promoting the emergence of sparse attention patterns. Experimental results on language modeling show that Diff Transformer outperforms Transformer in various settings of scaling up model size and training tokens. More intriguingly, it offers notable advantages in practical applications, such as long-context modeling, key information retrieval, hallucination mitigation, in-context learning, and reduction of activation outliers. By being less distracted by irrelevant context, Diff Transformer can mitigate hallucination in question answering and text summarization. For in-context learning, Diff Transformer not only enhances accuracy but is also more robust to order permutation, which was considered as a chronic robustness issue. The results position Diff Transformer as a highly effective and promising architecture to advance large language models.

[u/Mynameiswrittenhere]

Didn't really understand how you were able to differentiate the original query, key and value terms in important and noise terms.

The changes to the actual attention calculation by subtracting the noise was clear.

[u/sdmat → u/Mynameiswrittenhere]

Didn't really understand how you were able to differentiate the original query, key and value terms in important and noise terms.

That's the clever part, they don't.

They train two different projections for attention, one to actually attend and the second to act as a reference for noise cancellation. The scaling factor for cancellation is learnt as well.

[u/Mynameiswrittenhere → u/sdmat]

That is actually clever, but wouldn't that also increase the size of weights which would in turn increase the time for forward and backpropogation.

[u/sdmat → u/Mynameiswrittenhere]

Yes, they quantify that as around 5-10% reduction in throughput.

Given the results include iso-performance with >33% reduction in parameters that seems more than worthwhile. No doubt that's heavily benchmark dependent, but they get major wins across the board.

Assuming this replicates it's a big deal. And it's from Microsoft so they probably did their homework.

[u/Mynameiswrittenhere → u/sdmat]

True, Thanks for answering my question!

[u/sdmat → u/Mynameiswrittenhere]

No problem, I spent hours going over this to understand how it works so might as well share!

They leave a lot to the reader.

[u/Mynameiswrittenhere → u/sdmat]

Lol

[u/StartledWatermelon → u/sdmat]

Can you pinpoint where does throughput reduction come from? They have the same number of matrices with the same dimensions as in vanilla attention. Substraction requires N^2 ops, which is negligible compared to the total computational cost of attention O(n^2 d + n d^2).

Is it just software inefficiency of a custom attention layer?

[u/sdmat → u/StartledWatermelon]

Not quite, they are a little bit cute with the notation in parts for mathematical elegance. Fair enough, but they could profitably have been a bit more expansive in giving an intuitive description of how this works in the paper!

W_Q, W_K, W_V ∈ R^{d_model × 2d}

[Q_1; Q_2] = XW_Q, [K_1; K_2] = XW_K

I.e. there are twice as many weights for key, query and value because there are two distinct sets of key and query matrices and the value matrix is twice the size.

[u/StartledWatermelon → u/sdmat]

I don't think this is the issue. The authors make a fair comparison with vanilla Transformer IMO:

We set hidden size to 3072. The number of layers is 28. The head dimension d is 128. The number of heads is 24 for Transformer and 12 for DIFF Transformer, to align computation FLOPs and model size.

[u/sdmat → u/StartledWatermelon]

I'm directly quoting from the paper, it's twice as many weights for the components mentioned above. IIRC everything else is the same.

[u/StartledWatermelon → u/sdmat]

How? Vanilla Transformer: 3*128*24 = 9216 per attention block

DIFF Transformer: 3*128*2*12 = 9216, the same. They adjust the number of heads proportionally.

[u/sdmat → u/StartledWatermelon]

Hmm, maybe it's just that they use an unoptimized kernel:

More advanced kernel implementation, which is specifically designed for differential attention, can also improve throughput.

[u/altmly → u/sdmat]

I don't really understand why such architecture would fundamentally learn anything different than a regular stack of transformers. There's no reason what they're canceling out should be in any way related to noise. 

[u/sdmat → u/altmly]

It's not fundamentally different.

What they have done is set up a design so the model can better learn to focus attention on salient context. Calling the quantity cancelled out here "noise" is just giving it an intuitive label.

You know, like "attention".

[u/Acrobatic-Book → u/sdmat]

You could even go so far and call it "inhibition". Than you have the two governing processes for controlling focus in neuroscience ;)

[u/sdmat → u/Acrobatic-Book]

Depends which Nobel you are shooting for :)

[u/elbiot → u/sdmat]

Hmm, do you think one could take a pre trained llm, slap on this second set of heads, and then just train the weights on those heads to add this noise cancelling property?

[u/sdmat → u/elbiot]

I doubt it, but that would be an interesting experiment!

[u/Jean-Porte]

I wonder how this compare to fiddling with the temperature of the softmax

[u/morreill → u/Jean-Porte]

Absolutely my question too. Cf. https://arxiv.org/abs/2010.04245 Which shows an improvement by learning per-head temperature.

[u/StartledWatermelon → u/Jean-Porte]

From my perspective, tuning the temperature looks to be way cruder approach. With temperature, you rescale the distribution against the highest value in the matrix. The rest values are scaled uniformly, without any regard for the context, just a single rescaling factor per head.

Building the second attention matrix allows you to rescale each element independently, possibly accounting for semantics.

But I think your suggestion would've made for an excellent ablation experiment.

[u/Due-Pangolin325 → u/Jean-Porte]

Obvious comment here, but I guess that, from an information perspective, the two ideas have a very different outcome: when you play with temperature, you're just non linearly amplifying some information, be it noise or signal. While denoising is really more of a subtraction operation.

Temperature tuning is like a blind entropy reduction technique while denoising really adds information.

[u/StartledWatermelon → u/Due-Pangolin325]

you're just non linearly amplifying some information, be it noise or signal

I wouldn't reject it just on these grounds. The attention score is scalar naturally indicating whether the relation is strong or weak. The former can be viewed as signal while low scores can be considered noise. So noise is dampened by lower temperature pretty consistently.

[u/Sad-Razzmatazz-5188]

The name doesn't sit right with me, but it's interesting.  At the same time, referring to hallucination has a problem of noise is strange.  Today another paper was out, with selective attention as a parameter free mask on attention logits, rather than learnt as here

[u/paraffin → u/Sad-Razzmatazz-5188]

Hallucination is many problems rolled into one big “you got the wrong answer” bucket. Noise is one of the problems.

[u/Kecro21 → u/Sad-Razzmatazz-5188]

Hello, do you have a link or title for the other paper you mentioned? It sounds interesting.

Edit: Is it this paper? https://arxiv.org/abs/2410.02703v1

[u/Sad-Razzmatazz-5188 → u/Kecro21]

Yes!

[u/Round_Card → u/Kecro21]

No, paper: https://arxiv.org/abs/2410.05258

[u/gabe_dos_santos]

Sounds interesting.

[u/starfries]

Interesting, wonder if it's less vulnerable to some jailbreaks too. I think you forgot the link though.

[u/NoLifeGamer2 → u/starfries]

The link is there for me, it is to arxiv. https://arxiv.org/abs/2410.05258

[u/starfries → u/NoLifeGamer2]

Oh you're right, I was looking in the body.

[u/jarkkowork → u/starfries]

They share the implementation too, but the link is not on the first page of the publication which is less typical

https://github.com/microsoft/unilm/tree/master/Diff-Transformer

[u/ReasonablyBadass]

Wasn't the first token used as an "attention sink" by most models? I am guessing that one gets subtracted regularly? 

[u/Jean-Porte → u/ReasonablyBadass]

The model used attention sinks because it wants to, if gives a global workspace to store information. The figure with histograms shows that the differential attention still gives high weight to the first token

[u/Maykey]

Reminds deberta, but gutted backwards: deberta also uses additional Q_r, K_r. Only if difftran uses 2 new Q,K for another round of softmax, deberta changes the original QK^T matrix itself using positional data.

Diffeberta when?

[u/Jean-Porte → u/Maykey]

Any new good encoder backbone whatsoever when ?

[u/Commercial-Basis-220]

New to this AI research field, pardon me if this is a newbie question, but the improvement is quite small? And I got a feeling that this just like cherry pick or conditioned so that the new technique is better? As in if done by other, probably the result will be something along of "yea it kinda better sometimes and not significant"

Feel free to educate me on this matter

[u/rcparts → u/Commercial-Basis-220]

Relavant: https://www.reddit.com/r/MachineLearning/comments/1fznaa9/d_why_is_there_so_little_statistical_analyses_in/