Co je FlashAttention a jaký přínos má pro paměť?

FlashAttention je optimalizovaný algoritmus pro výpočet mechanismu pozornosti (attention) v transformerech, který výrazně snižuje spotřebu paměti a urychluje výpočty ve srovnání s klasickým přístupem.

Hlavní myšlenkou FlashAttention je vypočítat pozornost bez explicitního ukládání velkých matic mezivýsledků (například matice vah pozornosti QK^T), pomocí blokové zpracování a efektivního využití cache GPU. To umožňuje snížit maximální spotřebu paměti přibližně o 3-4krát, což je obzvlášť důležité při tréninku velkých jazykových modelů (LLM) s dlouhými sekvencemi.

Přibližně, pokud klasická pozornost vyžaduje paměť O(N^2) pro sekvenci délky N, FlashAttention tuto zátěž snižuje díky výpočtům v reálném čase a optimalizacím, což umožňuje spouštět modely s většími kontexty na stejných GPU.

Tímto způsobem je FlashAttention způsob, jak učinit pozornost efektivnější z hlediska paměti a rychlosti, což je kritické pro velké modely.