[feat] Compositional attention

[blefaudeux]

What does this PR do?

Implements Compositional Attention (based on the reference implementation), as mentioned in #41

Paper

TODOs

• Sane defaults
• Speedup wherever possible. Looks like it takes a lot of memory also at the moment, probably some dummy mistakes
• Maybe self-attention optimization (single proj) -> doable if moving the projections within the attention to the inproj class, worth it?
• Add a lot of explanations/documentations
• [0] Some IR results ? -> that would be for another task probably ?

cc @sarthmit if interested

Before submitting

• Did you have fun?
  • Make sure you had fun coding 🙃
• Did you read the contributor guideline?
• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  • N/A
• Did you make sure to update the docs?
  • N/A
• Did you write any new necessary tests?
  • N/A
• Did you update the changelog? (if needed)
  • N/A

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

[sarthmit]

Please feel free to reach out. I'd be happy to help :)

[blefaudeux]

Please feel free to reach out. I'd be happy to help :)

@sarthmit actually, I have a question: I've seen quite a few .contiguous().do_something() (for instance), it's not completely obvious to me that it's beneficial since you change the tensor structure right after that, and you don't seem to be using specific kernels which would require having a contiguous tensor in the first place (there are actually a bunch of kernels in xformers which would love that, but .contiguous would be last in that case, not first).
Is there another reason that I'm missing, TPUs for instance ?

[sarthmit]

@blefaudeux I think you are right, there is probably no use for those .contiguous(). I just used the fairseq multi-head codebase as the starting place and they seem to have it also, so I never bothered to remove it. See: multi-head

[codecov-commenter]

Codecov Report

Merging #178 (ca0a693) into main (c16078b) will increase coverage by 0.39%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main     #178      +/-   ##
==========================================
+ Coverage   90.62%   91.02%   +0.39%     
==========================================
  Files          58       59       +1     
  Lines        2892     3019     +127     
==========================================
+ Hits         2621     2748     +127     
  Misses        271      271              

Flag	Coverage Δ
Python	91.02% <100.00%> (+0.39%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
xformers/components/attention/favor.py	100.00% <ø> (ø)
xformers/components/__init__.py	100.00% <100.00%> (ø)
xformers/components/attention/attention_mask.py	98.50% <100.00%> (ø)
xformers/components/attention/base.py	97.22% <100.00%> (+0.16%)	⬆️
xformers/components/attention/compositional.py	100.00% <100.00%> (ø)
xformers/components/multi_head_dispatch.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update c16078b...ca0a693. Read the comment docs.

[blefaudeux]

microGPT test (causal, autoregressive)

Friends of my soul: set down.
What stands before I came to come?

BENVOLIO:
Madam, her were born.

MERCUTIO:
Tut, such more sorrow we did not so long.
I have no stand of small as a torment for a
point a man else.

ANTONIO:
Where is the brother of him?

SEBASTIAN:
Or despite of all the sun.

ANTONIO:
Go, marry, of many choice, my mother,
Would the near I let thee be freed:
As with, sir, a brave form'd her by my mood
Shall twenty years a traitor to turn: tell them
Show'd their misfer when seems but when I shall never
Their princely hate with the stealing death.

DUKE OF AUMERLE:
I do fear thee, my daughter not stay.

KING RICHARD II:
Say that's not swine may know thou art.

QUEEN ELIZABETH:
To thou art deceived: if thou both part to tell,
Thou wouldst be thing in she foul born in presence.
If I know not, Grumio that thought it good
Is mutually: here's a stale, if thou beholdst him
Are our bodies cousin with customary here,
A fine schoolmaster; there be we fined,
No, nor do as so fast; it is too long

But loss not as good as a simple attention with the defaults (it's not really the target for this attention mechanism though):

EDIT: This was due to the rotary embeddings being used for the MHA, and not here (not supported with this PR). Without the rotary embeddings there's not a lot of difference on this task, the perf loss disappears

[blefaudeux]

@sarthmit if you have cycles for a review, that would be great ! I've removed a few options (not that many actually) and aligned the terms with the other attentions here, + reused existing building blocks that we have (input projections for instance, there's an optimization to be used for self-attention

[blefaudeux]

reducing the memory load, generating a new graph for everyone

[blefaudeux]

hotfix, this was needed actually, not directly tied to this PR

[blefaudeux]

I tried to align the different field names and reformulate wherever possible, but there are still some specificities. It's more visible here since it's a specific attention unit test, but for real the new fields were already needed for the MHA and they don't need to be duplicated

[blefaudeux]

code coverage, this attention exposes quite a few knobs which are mostly orthogonal to the other attentions, so I figured it was best to cover them in a dedicated unit test

[blefaudeux]

compositional takes more memory, so this task does not pass on a 8GB gpu (desktop 3080) without bumping up the accumulation

[blefaudeux]

convenience, remove the need for field duplication in between attention and MHA

[blefaudeux]

it's already doing that by default I believe, just making it a more explicit

[blefaudeux]

difference with the original implementation: reuse xformers input projection, because the task is the same, and there are some optimizations to be had here (if self attention, project once to get q,k,v instead of three calls for instance)

[dianaml0]

Interesting approach!

[dianaml0]

Could we just use the softmax here?

[blefaudeux]

good catch, fixing that !

[sarthmit]

@sarthmit if you have cycles for a review, that would be great ! I've removed a few options (not that many actually) and aligned the terms with the other attentions here, + reused existing building blocks that we have (input projections for instance, there's an optimization to be used for self-attention

@blefaudeux What do you mean by cycles for a review?

By the way, there is also an easy way to test/debug the code. If you set the Value Scores in Equation 13 to the Identity matrix (here), then you should recover the standard Multi-Head Attention performance.

[blefaudeux]

@sarthmit if you have cycles for a review, that would be great ! I've removed a few options (not that many actually) and aligned the terms with the other attentions here, + reused existing building blocks that we have (input projections for instance, there's an optimization to be used for self-attention

@blefaudeux What do you mean by cycles for a review?

By the way, there is also an easy way to test/debug the code. If you set the Value Scores in Equation 13 to the Identity matrix (here), then you should recover the standard Multi-Head Attention performance.

I meant "brain cycles" :) (aka free time). I can have a look later today, but the logic should be maintained

[blefaudeux]

@sarthmit if you have cycles for a review, that would be great ! I've removed a few options (not that many actually) and aligned the terms with the other attentions here, + reused existing building blocks that we have (input projections for instance, there's an optimization to be used for self-attention

@blefaudeux What do you mean by cycles for a review?

By the way, there is also an easy way to test/debug the code. If you set the Value Scores in Equation 13 to the Identity matrix (here), then you should recover the standard Multi-Head Attention performance.

it's a good idea I think, testing that right now, I'm still seeing a difference in perf, investigating. I can see that there are other differences with the Vaswani paper, for instance the normalization is not the same (/sqrt(dim_head) vs /sqrt(dim model)), but I don't think that's enough to explain the difference

[blefaudeux]

ok, confirmed @sarthmit @dianaml0 when forcing the values score to be identity the perf are in line with MHA: the difference in the above is due to.. the rotary embeddings ! These are not part of this PR but they were used in the comparison above for the classical mechanism, and they explained the gap (without them the two mechanisms are aligned if the value score is identity, as expected). All good for landing I think, and it means that supporting rotary embeddings in compositional attention could be a follow up

[sarthmit]

@blefaudeux Awesome! I can try sparing some time if needed, is there anything in particular you want me to look at?

I would also be curious to see if you see performance gains on any of the tasks (without the identity matrix)

[blefaudeux]

@blefaudeux Awesome! I can try sparing some time if needed, is there anything in particular you want me to look at?

I would also be curious to see if you see performance gains on any of the tasks (without the identity matrix)

I've not checked LRA but that could be interesting to get the scores, even if I think that LRA could do with more and better tasks (see this paper from @xwhan for instance https://arxiv.org/pdf/2112.07210.pdf). Rotary embeddings also seem to make a big difference with NLP, would be nice to support them with compositional, should not be a big change.

There's a definite perf impact though (at least with our implementations, even when pulling in some semi-optimized parts from xformers), it could be another axis to look at because I think that a lot of folks will "just" bump up the number of heads with the scaled dot product instead if compositional has too much of a speed/memory impact

[sarthmit]

@blefaudeux Awesome! I can try sparing some time if needed, is there anything in particular you want me to look at?
I would also be curious to see if you see performance gains on any of the tasks (without the identity matrix)

I've not checked LRA but that could be interesting to get the scores, even if I think that LRA could do with more and better tasks (see this paper from @xwhan for instance https://arxiv.org/pdf/2112.07210.pdf). Rotary embeddings also seem to make a big difference with NLP, would be nice to support them with compositional, should not be a big change.

There's a definite perf impact though (at least with our implementations, even when pulling in some semi-optimized parts from xformers), it could be another axis to look at because I think that a lot of folks will "just" bump up the number of heads with the scaled dot product instead if compositional has too much of a speed/memory impact

Thanks, I will check the paper out.

How big is the perf impact? In favour of MHA or in favour of compositional?

[blefaudeux]

@blefaudeux Awesome! I can try sparing some time if needed, is there anything in particular you want me to look at?
I would also be curious to see if you see performance gains on any of the tasks (without the identity matrix)

I've not checked LRA but that could be interesting to get the scores, even if I think that LRA could do with more and better tasks (see this paper from @xwhan for instance https://arxiv.org/pdf/2112.07210.pdf). Rotary embeddings also seem to make a big difference with NLP, would be nice to support them with compositional, should not be a big change.

There's a definite perf impact though (at least with our implementations, even when pulling in some semi-optimized parts from xformers), it could be another axis to look at because I think that a lot of folks will "just" bump up the number of heads with the scaled dot product instead if compositional has too much of a speed/memory impact

Thanks, I will check the paper out.

How big is the perf impact? In favour of MHA or in favour of compositional?

By perf I meant compute time impact, sorry that wasn't clear. Compositional (this PR at least) is significantly slower than MHA (from this repo again) for the same number of heads, even if I understand that it captures way more relations. Could be worth a second eye on whether there are low hanging fruits?