rishi

I like this paper late last year which proposes a more compute-efficient and potentially more performative variant of multi-head latent attention. They call it compressed convolutional attention (CCA). There are a few inspiring takeaways from it. - What they did: MLA compresses QKV to a low-rank latent space, but it's mainly for KV caching. QKV is immediately projected back after storage is done. The authors instead propose to do not only caching but also computation in the latent space. Concretely this means after low-rank projection, the whole softmax attention is done there, then projected back to hidden dim (p1). Intuitively, this will not work well because low-rank projection loses a lot of information. So the authors apply a token-wise convolution and then a head-dim-wise convolution to QK. Ablations show that the two convolutions contribute a lot to improved loss and evaluation results. As is in my pseudocode (p2), they make a bunch of other interesting configurations, such as concating two of the most recent Vs, and adding a "QK mean" to Q and K to mix them together. And this can be made even more efficient if we add grouped query attention on top, ie. basically expand KV by num_groups so that we use even fewer parameters. - How it performs: the authors did some detailed evaluations against MHA, MLA, and GQA. They make sure to match the total params count, and for MLA and GQA they also match the KV cache compression rate. They show that the GQA-enhanced CCA perform on par with the lossless MHA on several benchmarks (HellaSwag, Winogrande, the usual stuff) as well as the final training loss. They note that this is especially handy for MoE models, because a smaller attention module means a bigger MoE module and specially each expert with a bigger size. Overall, this is very cool work, and belongs to the line of research that lets attention stay at quadratic but just keeps its activation size in check. It'll be interesting to calculate a FLOP breakdown between this (especially the GQA-enhanced one) and the new Deepseek Sparse Attention, since DSA is by design lossless and potentially also very efficient.

How to ("properly") orthogonalize convolutional layers for Muon optimizer Trick: Assume circular kernels to allow diagonalization Blog post + proof of concept CIFAR10 speedrun fork (unoptimized and slow for now but better convergence per step)

it's still climbing higher. 459B input and 2.6B output, 176:1.🤔