LongNet: Scaling Transformers to 1,000,000,000 Tokens

This works propose a dilated attention which expands the attentive field exponentially as the distance between the tokens grow.

Let’s first see an example to make the idea intuitive lets say we are calculating the attention between tokens within a book, lets consider the first, segment length 4 as a single paragraph, the next segment 8 as a chapter and the last with segment length 16 as the full book.

As you might expect the tokens existing in a single attention will have a dense relationship where each token attending ever other token in the space. Now lets see the tokens in a paragraph space, they will attend the related tokens but since the distance between the tokens increase it becomes sparce. Similarly as we expand from the chapter to the whole book, even the book keeps a theme which makes the tokens some what related but the relation between tokens in chapter k and chapter z will be sparce.

Note that no matter how much space you’re covering the number of dot products you’re performing will be same in our case 10 dot products. This will let you expand over a very large token space while keeping the localized attention and keeping a sparce attention between tokens as the distance between them grows.

Note that that, each block is independent of the other and we will be calculating a localized attention. Then we concatenate them and calculate the weighted average.

When we have multiple heads, we can jump in j-1 gaps in each head to do sparse and efficient computation. For example lets say we have, 4 heads with segment 8 and dilation rate of 2. First we start from 0,0 and go 2,0 and 2,2 and then 4,0, 4,2 and 4,4 and then 6,0, 6,2, 6,4 and 6,8. Then on the other head we do 1,1, then 3,1 and 3,3, 5,1 5,3 and 5,5 and 7,1 7, 3, 7,5 and 7,7.

One interesting factor in this paper is the attention propagation, meaning even two tokens are not related directly through dot product, we can still find a path from one to another in magnitude of log n.

For example, between 0 and 15th token there is no direct product but we can use the path 0,12 → 12,12 → 12,15 → 15,15 to reach one to another.

This means if we expand our expand our token space by 10 we are only adding 1 more step.

The other interesting thing is the computation complexity, in standard vanilla we will be facing N * v^2 * d but in this work since we divided our vector space into w windows which dilated at the rate of r our flops operation will be

Which can also be dilated across multiple segment rate giving us

which yield us a complexity that grows linearly with O(Nd)

Similarly, the independency of the blocks enable us to do a distributed training.

References

Original Paper: Ding, Jiayu, et al. “Longnet: Scaling transformers to 1,000,000,000 tokens.” arXiv preprint arXiv.02486 (2023).