Τι είναι το FlashAttention και ποιο είναι το κέρδος στη μνήμη;

FlashAttention είναι ένας βελτιστοποιημένος αλγόριθμος για τον υπολογισμό του μηχανισμού προσοχής (attention) στους μετασχηματιστές, ο οποίος μειώνει σημαντικά την κατανάλωση μνήμης και επιταχύνει τους υπολογισμούς σε σύγκριση με την κλασική προσέγγιση.

Η κύρια ιδέα του FlashAttention είναι να υπολογίζει την προσοχή χωρίς να αποθηκεύει ρητά μεγάλες μήτρες ενδιάμεσων αποτελεσμάτων (π.χ., τον πίνακα βαρών προσοχής QK^T), χρησιμοποιώντας επεξεργασία σε μπλοκ και αποτελεσματική χρήση της cache της GPU. Αυτό επιτρέπει τη μείωση της μέγιστης κατανάλωσης μνήμης περίπου 3-4 φορές, κάτι που είναι ιδιαίτερα σημαντικό κατά την εκπαίδευση μεγάλων μοντέλων γλώσσας (LLM) με μακριές ακολουθίες.

Περίπου, αν η κλασική προσοχή απαιτεί μνήμη O(N^2) για μια ακολουθία μήκους N, το FlashAttention μειώνει αυτό το φορτίο μέσω υπολογισμών σε πραγματικό χρόνο και βελτιστοποιήσεων, επιτρέποντας την εκτέλεση μοντέλων με μεγαλύτερα συμφραζόμενα στις ίδιες GPU.

Έτσι, το FlashAttention είναι ένας τρόπος να κάνουμε την προσοχή πιο αποδοτική σε μνήμη και ταχύτητα, κάτι που είναι κρίσιμο για μεγάλης κλίμακας μοντέλα.