@inproceedings{
singhania2024loki,
title={Loki: Low-rank Keys for Efficient Sparse Attention},
author={Prajwal Singhania and Siddharth Singh and Shwai He and Soheil Feizi and Abhinav Bhatele},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=raABeiV71j}
}
Inference on large language models (LLMs) can be expensive in terms of the
compute and memory costs involved, especially when long sequence lengths are
used. In particular, the self-attention mechanism used in LLM inference contributes
significantly to these costs, which has sparked an interest in approximating the self-attention 
computation to reduce such costs. In this work, we propose to approximate
self-attention by focusing on the dimensionality of key vectors computed in the
attention block. Our analysis reveals that key vectors lie in a significantly lower-dimensional 
space, consistently across several datasets and models. Exploiting this
observation, we propose Loki, a novel sparse attention method that ranks and selects
tokens in the KV-cache based on attention scores computed in low-dimensional
space. Our evaluations show that Loki is able to speed up the attention computation
due to reduced data movement (load/store) and compute costs while maintaining
the efficacy of the models better than other popular approximation methods.
Loki: Low-rank Keys for Efficient Sparse Attention
Prajwal Singhania Siddharth Singh Shwai He Soheil Feizi Abhinav Bhatele