Block Sparse Attention with Log-Linear Complexity

Les chercheurs proposent PISA, un mécanisme d'attention parcimonieuse par blocs qui utilise une stratégie de sélection pyramidale Top-K pour réduire la complexité de calcul de quadratique à log-linéaire. PISA construit une hiérarchie de clés du grossier au fin par regroupement et applique un score LogSumExp sur des ensembles de candidats bornés. Ils développent des noyaux Triton conscients du matériel pour l'entraînement et l'inférence qui fusionnent le routage hiérarchique et le calcul de score sans matérialiser la matrice de score requête-clé, obtenant des résultats comparables sur le raisonnement de bon sens et de meilleures performances sur les tâches de recherche par rapport aux modèles de référence.

Through pooling, we construct O(log N) levels of keys, yielding an overall complexity of O(Nlog N), where N denotes the sequence length.