Čo je FlashAttention a aký prínos má pre pamäť?

FlashAttention je optimalizovaný algoritmus na výpočet mechanizmu pozornosti (attention) v transformeroch, ktorý výrazne znižuje spotrebu pamäte a zrýchľuje výpočty v porovnaní s klasickým prístupom.

Hlavná myšlienka FlashAttention je vypočítať pozornosť bez explicitného ukladania veľkých matíc medzi výsledkami (napríklad matice váh pozornosti QK^T), pomocou blokového spracovania a efektívneho využívania cache GPU. To umožňuje znížiť maximálnu spotrebu pamäte približne o 3-4 krát, čo je obzvlášť dôležité pri tréningu veľkých jazykových modelov (LLM) s dlhými sekvenciami.

Približne, ak klasická pozornosť vyžaduje pamäť O(N^2) pre sekvenciu dĺžky N, FlashAttention túto záťaž znižuje pomocou výpočtov v reálnom čase a optimalizácií, čo umožňuje spúšťať modely s väčšími kontextami na tých istých GPU.

Týmto spôsobom je FlashAttention spôsob, ako spraviť pozornosť efektívnejšou z hľadiska pamäte a rýchlosti, čo je kritické pre veľké modely.