4 ms·
This looks really interesting! If these guys succeed in bringing self-attention's computational cost down from O(n²) to O(n log n), that would be a huge win. Th
by cs702 4y ago
This looks really interesting! If these guys succeed in bringing self-attention's computational cost down from O(n²) to O(n log n), that would be a huge win. The quadratic cost makes it very difficult to increase sequence length on current hardware. I'm going to take a closer look.
There are other interesting ongoing efforts to increase sequence length. One that has worked for me is this dynamic routing algorithm, related to self-attention, that can handle sequences with 1M+ tokens in a single GPU: https://github.com/glassroom/heinsen_routing https://github.com/glassroom/heinsen_routing . Right now, you can take 1,000 sequences of hidden states computed by a pretrained transformer, each sequence with, say, 1024 tokens, concatenate them into a single ultra-long sequence with 1,024,000 hidden states, slap 1,024,000 position encodings on top, and feed the whole thing to that routing algorithm to predict the next token (or whatever other training objective you want to optimize for). It works. Search the README for "Very Long Sequences".
If anyone here has other suggestions for working with long sequences (hundreds of thousands to millions of tokens), I'd love to learn about them.
- famouswaffles 4y agoThere are already linear attention advances. Gpt-4-32k is almost certainly using some form of flash attention. Attention isn't really O(n2) anymore.
- cs702 4y agoMy understanding is that FlashAttention's memory use is linear, or close to linear in practice, but computation is still O(n²). I'm unaware of anyone being able to apply FlashAttention on, say, a million tokens, because it must execute ~1/2 x 1,000,000^2 x n_head dot-products, each in a subspace with d_head dimensions. That's not exactly computationally cheap!
- famouswaffles 4y agoNo you're right. I mistook you. Compute isn't linear yet.
- lucidrains 4y agoIt is only linear in terms of memory, not compute. Flash attention is a big advance, but not enough for 1 million tokens
- deleted 4y ago[deleted]
- deleted 4y ago[deleted]
- deleted 4y ago[deleted]
- logicchains 4y agohttps://github.com/BlinkDL/RWKV-LM https://github.com/BlinkDL/RWKV-LM this claims to work well with long sequences.
- cs702 4y agoThank you. The evidence I've seen suggests RWKV-LM does not work well with long sequences: https://news.ycombinator.com/item?id=35375616 https://news.ycombinator.com/item?id=35375616 So far, every method I've come across that works well with long sequences does it by explicitly modeling (approximating) all possible pairwise relationships between tokens. In the case of attention, that means modeling n×n, or n², pairwise relationships. In the case of that dynamic routing method, it means modeling n×m pairwise relationships, where n and m are the number of tokens in the input and output sequences. I suspect that the key reason why that routing method works for long sentences is because m can be much smaller than n: Routing 1M tokens to 100 tokens requires ~10,000x less compute than applying attention over 1M tokens for each of 1M tokens. If the work described by the OP succeeds in bringing computational cost down from O(n²) to O(n log n) while still explicitly modeling (approximating) all n² pairwise relationships, that would be a huge win.
- smaddox 4y agoOther approaches for long sequence lengths: - "Hyena Hierarchy: Towards Larger Convolutional Language Models" https://arxiv.org/pdf/2302.10866.pdf https://arxiv.org/pdf/2302.10866.pdf - "Resurrecting Recurrent Neural Networks for Long Sequences" https://arxiv.org/pdf/2303.06349.pdf https://arxiv.org/pdf/2303.06349.pdf It's unclear (at least to me) if any of these scale as well as transformers at this point.
- cs702 4y agoThank you. The OP is from the authors of Hyena, and I agree it's unclear if it will scale as well as self-attention in transformers. The paper on RNNs looks very interesting. I like the fact that it keeps things simple and familiar. I've added it to my reading list.
- smaddox 4y agoYeah, just saw that. Sorry for the noise.
- cs702 4y agoNot a big deal. The paper on RNNs is exactly the sort of thing I was hoping someone would bring up here. Thank you!
- ilaksh 4y agoHow does that one compare to RWKV?
- smaddox 4y agoRWKV is random chance on the PathX problem. Both of those models are > 95% on it, if I remember correctly. I don't think Hyena can be run in RNN mode, though. But it _might_ scale better than the LRU model (second paper).
- smaddox 4y agoThis EfficientVectorRouting paper (and related papers) is fascinating. It seems to have endless potential applications. I'm surprised I had not heard of it before. I'm excited to experiment with novel model structures incorporating it. What's the catch? I'm guessing they're fairly expensive layers, since they're based on EM? Although it looks like it's using incremental EM (iEM), which perhaps makes it less expensive? May I ask what you are using it for?
- cs702 4y agoThere's no catch, really, but like everything else, there are applications and uses for which there are better tools. For example, if you need a layer for modeling n² relationships, you'd be much better off using something like FlashAttention than this routing method. I've found it to be a useful tool. It's become my go-to solution whenever I need to combine data, hidden states, or model outputs of different sizes/dimensions/modalities into something that works on the first try. It can "glue models together into systems" without much hassle. I mean, you still have to tinker and tweak to get things to work well, but no more than usual. Most recently, I've been using it to build models with access to long-term memory, more or less like this: Concatenate a large number of sequences of previously computed transformer hidden states. Apply a couple of routing layers to shrink them into a single short sequence of fixed size. Attach the shrunk sequence as a form of long-term memory to the current sequence of final transformer hidden states. Slap a couple of routing layers on top to predict things of interest. Finally, train only the routing layers (leave all the transformer stuff frozen). You can of course save the long sequences of previous transformer hidden states in a key-value store and fetch them as needed :-)
- ilaksh 4y agoThis is a really dumb question but could EfficientVectorRouting be used for text completion somehow?
- cs702 4y agoIn principle, yes, you can, by using something like this on the last layer https://github.com/glassroom/heinsen_routing#sequence-to-vector https://github.com/glassroom/heinsen_routing#sequence-to-vec... ... but keep in mind it will not match the training efficiency of triangular masking in autoregressive transformers. You'd have to train the model one time-step at a time, vs. all time-steps in parallel with a masked autoregressive transformer. If you're looking to make context longer, a more practical approach is to route multiple past sequences of frozen pretrained transformer hidden states to a much shorter sequence, say, of length C, concatenate that with the current final sequence of hidden states, say, of length N, apply masked autoregressive attention on the C+N hidden states and then a feedforward layer to predict all next tokens in parallel, and last, drop the first C predictions, i.e., use only the trailing N predicted next tokens for training and inference.