Πώς υλοποιείται η cache KV στην inference ενός αυτοπαραγωγικού μετασχηματιστή;

Στα αυτοαναφορικά μετασχηματιστικά μοντέλα, κατά τη διάρκεια της επεξεργασίας, είναι σημαντικό να χρησιμοποιούνται οι υπολογισμοί αποδοτικά ώστε να μην χρειάζεται να επαναυπολογίζεται ολόκληρη η ακολουθία από την αρχή. Για αυτό το σκοπό, χρησιμοποιείται η μνήμη cache KV, η οποία αποθηκεύει τα κλειδιά (Key) και τις τιμές (Value) από τα επίπεδα αυτοπροσοχής.

Η ιδέα είναι ότι, κατά τη δημιουργία του επόμενου token, έχουμε ήδη υπολογίσει τα κλειδιά και τις τιμές για όλα τα προηγούμενα tokens. Αντί να τα υπολογίσουμε ξανά, τα αποθηκεύουμε στη μνήμη (cache) και, κατά τον υπολογισμό της προσοχής για το νέο token, χρησιμοποιούμε αυτήν την cache.

Αυτό λειτουργεί περίπου ως εξής:

• Για κάθε επίπεδο του μετασχηματιστή, κατά την επεξεργασία ενός token, τα κλειδιά και οι τιμές αποθηκεύονται σε ξεχωριστές μπλοκ.
• Κατά τη δημιουργία του επόμενου token, υπολογίζονται μόνο τα κλειδιά και οι τιμές για αυτό, και στη συνέχεια συνδυάζονται με αυτά που έχουν ήδη αποθηκευτεί.
• Ο μηχανισμός προσοχής χρησιμοποιεί το συνδυασμένο σύνολο κλειδιών και τιμών για να υπολογίσει το πλαίσιο.

Αυτό επιταχύνει σημαντικά την επεξεργασία, μειώνοντας τη μαθηματική πολυπλοκότητα από τετράγωνη σε γραμμική ως προς το μήκος της ακολουθίας.

Παραδείγματα ψευδοκώδικα για ένα επίπεδο:

# kv_cache αποθηκεύει τα κλειδιά και τις τιμές των προηγούμενων tokens
new_key, new_value = compute_kv(new_token)
kv_cache.keys = concatenate(kv_cache.keys, new_key)
kv_cache.values = concatenate(kv_cache.values, new_value)
output = attention(query=new_key, keys=kv_cache.keys, values=kv_cache.values)

Έτσι, η cache KV είναι ένας μηχανισμός αποθήκευσης ενδιάμεσων κλειδιών και τιμών για την επιτάχυνση της διαδοχικής παραγωγής σε μετασχηματιστές.