Mamba: Ανακαθορίζοντας το Μοντέλο Σειράς και Υπερβαίνοντας την Αρχιτεκτονική Transformers

Στο άρθρο αυτό εξετάζουμε πώς το Mamba, ένα μοντέλο χώρου καταστάσεων (SSM) των Albert Gu και Tri Dao, αλλάζει τη μοντελοποίηση ακολουθιών. Επεξεργάζεται αποτελεσματικά σύνθετες ακολουθίες στη γλώσσα, τη γονιδιωματική και την ανάλυση ήχου. Οι επιλεκτικοί χώροι καταστάσεων επιτρέπουν υπολογιστικό κόστος γραμμικό ως προς το μήκος της ακολουθίας.

Το Mamba αντιμετωπίζει τις υπολογιστικές δυσκολίες των παραδοσιακών Transformer στις μεγάλες ακολουθίες. Η επιλεκτική προσέγγιση επιτρέπει ταχύτερη εξαγωγή συμπερασμάτων και γραμμική κλιμάκωση, βελτιώνοντας τον ρυθμό επεξεργασίας. Συνδυάζει γρήγορη επεξεργασία, επιλεκτικό επίπεδο SSM και σχεδιασμό προσαρμοσμένο στο υλικό, εμπνευσμένο από το FlashAttention. Σύμφωνα με τα αποτελέσματα που παρουσιάζει το αρχικό άρθρο, ξεπερνά πολλά υπάρχοντα μοντέλα, συμπεριλαμβανομένων μοντέλων βασισμένων σε Transformer.

Transformer και Mamba

Μοντέλα όπως το GPT-4 έχουν θέσει υψηλά πρότυπα στην επεξεργασία φυσικής γλώσσας, αλλά το κόστος επεξεργασίας αυξάνεται στις μεγάλες ακολουθίες. Το Mamba προτείνει μια διαφορετική αρχιτεκτονική για αποτελεσματικότερη επεξεργασία τους.

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

Η προσοχή υπολογίζεται από ερωτήματα, κλειδιά και τιμές που προκύπτουν από την είσοδο. Η σύγκριση των στοιχείων δίνει βάρη που εκφράζουν τη σχετική σημασία τους για την πρόβλεψη. Στην αρχιτεκτονική κωδικοποιητή–αποκωδικοποιητή που περιγράφεται εδώ, ο κωδικοποιητής επεξεργάζεται την είσοδο και ο αποκωδικοποιητής παράγει την έξοδο. Η περιγραφή αυτή αφορά τη συγκεκριμένη διάταξη Transformer, όχι κάθε μοντέλο της οικογένειας.

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

Το Mamba χρησιμοποιεί επιλεκτικούς χώρους καταστάσεων αντί για τον ίδιο μηχανισμό προσοχής. Σημείωση για την πηγή: η αντίστοιχη συγκριτική πρόταση στο αγγλικό άρθρο διακόπτεται στη μέση μιας λέξης· δεν αποδίδεται εδώ κάποιο άγνωστο συμπέρασμα.

Σε γενικές γραμμές, η λειτουργία ενός Transformer περιλαμβάνει:

• Επεξεργασία εισόδου: αναπαραστάσεις διανυσμάτων και πληροφορία θέσης μετατρέπουν την είσοδο σε μορφή κατάλληλη για το μοντέλο.
• Προσοχή: υπολογίζονται βαθμολογίες που καθορίζουν πόσο βάρος δίνεται στις άλλες θέσεις.
• Κωδικοποιητή και αποκωδικοποιητή: στη συγκεκριμένη διάταξη, πολλαπλά επίπεδα επεξεργάζονται την είσοδο και παράγουν την έξοδο.
• Προσοχή πολλαπλών κεφαλών: διαφορετικοί χώροι αναπαράστασης επιτρέπουν ταυτόχρονη εξέταση διαφορετικών σχέσεων και συμφραζομένων.
• Δίκτυα πρόσθιας τροφοδότησης ανά θέση: η ίδια μετατροπή εφαρμόζεται ανεξάρτητα σε κάθε θέση, με υπολειμματική σύνδεση και κανονικοποίηση επιπέδου.
• Παραγωγή εξόδου: ο αποκωδικοποιητής προβλέπει την επόμενη έξοδο με βάση το πλαίσιο του κωδικοποιητή και όσα έχει ήδη παραγάγει.

Η παράλληλη επεξεργασία και η προσοχή καθιστούν τους Transformer αποτελεσματικούς σε εργασίες όπως η μετάφραση και η παραγωγή κειμένου. Το Mamba αντιμετωπίζει το αυξανόμενο κόστος στις μεγάλες ακολουθίες μέσω επιλεκτικών χώρων καταστάσεων, με ταχύτερη εξαγωγή συμπερασμάτων και γραμμική κλιμάκωση ως προς το μήκος.

Mamba

Το Mamba απομακρύνεται από τα παραδοσιακά μπλοκ προσοχής και πολυεπίπεδων perceptron (MLP), δημιουργώντας μια ελαφρύτερη, ταχύτερη και περισσότερο ομοιογενή αρχιτεκτονική.

• Επιλεκτικά SSM: φιλτράρουν άσχετες πληροφορίες και εστιάζουν στα σχετικά δεδομένα, υποστηρίζοντας συλλογιστική που εξαρτάται από το περιεχόμενο.
• Αλγόριθμος προσαρμοσμένος στο υλικό: η παράλληλη επεξεργασία είναι βελτιστοποιημένη για σύγχρονο υλικό, ιδιαίτερα GPU, με ταχύτερους υπολογισμούς και μικρότερες απαιτήσεις μνήμης.
• Απλοποιημένη αρχιτεκτονική: η ενσωμάτωση επιλεκτικών SSM και η αφαίρεση των συμβατικών μπλοκ προσοχής και MLP ενισχύουν την κλιμάκωση και την απόδοση.

Τα αποτελέσματα που παρουσιάζονται καλύπτουν γλώσσα, ήχο και γονιδιωματική, τόσο στην προεκπαίδευση όσο και σε ειδικές εργασίες. Στη μοντελοποίηση γλώσσας, το Mamba φτάνει ή υπερβαίνει την απόδοση μεγαλύτερων Transformer στις αναφερόμενες συγκρίσεις. Ο κώδικας και τα προεκπαιδευμένα μοντέλα διατίθενται δημόσια στο GitHub.

Οι συνηθισμένες εργασίες αντιγραφής είναι απλές για γραμμικά μοντέλα. Η επιλεκτική αντιγραφή και οι κεφαλές επαγωγής απαιτούν δυναμική μνήμη που λαμβάνει υπόψη το περιεχόμενο.

Τα δομημένα μοντέλα χώρου καταστάσεων S4 συνδυάζουν χαρακτηριστικά αναδρομικών δικτύων (RNN), συνελικτικών δικτύων (CNN) και κλασικών μοντέλων χώρου καταστάσεων. Εμπνέονται από συνεχή συστήματα που αντιστοιχίζουν μονοδιάστατες συναρτήσεις ή ακολουθίες μέσω μιας λανθάνουσας κατάστασης. Προσφέρουν έναν αποτελεσματικό και προσαρμόσιμο τρόπο κατασκευής μοντέλων ακολουθιών.

Η δυναμική των μοντέλων S4

Ένα δομημένο SSM λαμβάνει την ακολουθία x και παράγει την έξοδο y, χρησιμοποιώντας μαθημένες παραμέτρους A, B, C και βήμα χρόνου Δ. Οι συνεχείς παράμετροι διακριτοποιούνται και κατόπιν εφαρμόζεται η λειτουργία SSM. Στη βασική χρονικά αμετάβλητη μορφή, οι παράμετροι παραμένουν ίδιες σε όλα τα βήματα χρόνου.

Η σημασία της διακριτοποίησης

Η διακριτοποίηση μετατρέπει συνεχείς παραμέτρους σε διακριτές με καθορισμένους τύπους, διατηρώντας τη σύνδεση με συστήματα συνεχούς χρόνου. Παρέχει ιδιότητες όπως αμεταβλητότητα ως προς την ανάλυση και κατάλληλη κανονικοποίηση, υποστηρίζοντας τη σταθερότητα και την απόδοση. Συνδέεται επίσης με τους μηχανισμούς πυλών των RNN, οι οποίοι ελέγχουν τη ροή πληροφορίας.

Γραμμικότητα και χρονική αμεταβλητότητα (LTI)

Στα S4, η δυναμική παραμένει σταθερή στον χρόνο και οι παράμετροι είναι ίδιες σε κάθε βήμα. Αυτή η ιδιότητα στηρίζει την αναδρομή και τη συνέλιξη, προσφέροντας ένα απλό αλλά ισχυρό πλαίσιο μοντελοποίησης ακολουθιών.

Υπέρβαση των βασικών περιορισμών

Η χρονική αμεταβλητότητα περιορίζει την προσαρμογή σε δεδομένα που απαιτούν μεταβαλλόμενη δυναμική. Η ερευνητική εργασία εισάγει παραμέτρους που μεταβάλλονται στον χρόνο, αίροντας τον περιορισμό LTI και διευρύνοντας τις ακολουθίες και τις εργασίες που μπορούν να αντιμετωπιστούν.

Ο όρος «μοντέλο χώρου καταστάσεων» χρησιμοποιείται ευρύτερα για αναδρομικές διαδικασίες με λανθάνουσα κατάσταση σε πολλούς επιστημονικούς κλάδους. Στη βαθιά μάθηση, τα S4 ή δομημένα SSM είναι συγκεκριμένη κατηγορία βελτιστοποιημένη για αποδοτικούς υπολογισμούς και σύνθετες ακολουθίες. Ενσωματώνονται σε νευρωνικά δίκτυα από άκρο σε άκρο ως μετασχηματισμοί ακολουθιών, με ρόλο ανάλογο των συνελικτικών επιπέδων στα CNN.

SSM και SSM με επιλογή

Γιατί χρειάζεται επιλεκτικότητα στη μοντελοποίηση ακολουθιών

Δομημένα SSM

Μια βασική λειτουργία είναι η συμπίεση του πλαισίου σε διαχειρίσιμη κατάσταση. Η επιλεκτική εστίαση ή το φιλτράρισμα της εισόδου βοηθούν στη διατήρηση χρήσιμης πληροφορίας μέσα σε αυτή τη συμπιεσμένη αναπαράσταση. Έτσι το μοντέλο ελέγχει προσαρμοστικά τη ροή πληροφορίας κατά μήκος της ακολουθίας, κάτι σημαντικό για σύνθετες γλωσσικές και άλλες εργασίες.

Στα επιλεκτικά SSM, οι παράμετροι εξαρτώνται από την είσοδο. Η δυναμική γίνεται χρονικά μεταβαλλόμενη και δεν μπορεί πλέον να χρησιμοποιήσει τη συνήθη αποδοτική συνέλιξη των χρονικά αμετάβλητων μοντέλων. Αντί γι’ αυτήν χρησιμοποιείται αναδρομική επεξεργασία με γραμμική κλιμάκωση.

Στην παραλλαγή SSM με επιλογή (S6), οι παράμετροι B, C και Δ εξαρτώνται από την είσοδο x. Η διακριτοποίηση λαμβάνει υπόψη αυτή την επιλογή και μια λειτουργία σάρωσης επεξεργάζεται την ακολουθία, προσαρμόζοντας δυναμικά την εστίαση σε κάθε βήμα.

Σημαντικά αποτελέσματα απόδοσης του Mamba

Το διάγραμμα της πηγής παρουσιάζει το Mamba ως καλύτερο στην κατηγορία του σε κάθε εικονιζόμενη αξιολόγηση.

Το Mamba παρουσιάζει υψηλή ταχύτητα εξαγωγής συμπερασμάτων και ακρίβεια στις αναφερόμενες δοκιμές. Αξιοποιεί μεγαλύτερα συμφραζόμενα στη μοντελοποίηση DNA και ήχου και υπερβαίνει προηγούμενα μοντέλα σε σύνθετες εργασίες που απαιτούν εξαρτήσεις μεγάλης εμβέλειας. Οι αξιολογήσεις χωρίς παραδείγματα εκπαίδευσης για τη συγκεκριμένη εργασία (zero-shot) αναδεικνύουν επίσης την ευελιξία, την αποδοτικότητα και την κλιμάκωσή του.

Πρώτα βήματα με το Mamba

Οι τεχνικές απαιτήσεις που παραθέτει το άρθρο του 2023 είναι Linux, GPU NVIDIA, PyTorch 1.12+ και CUDA 11.6+. Η εγκατάσταση γίνεται με pip από το αποθετήριο Mamba. Για προβλήματα συμβατότητας με εκδόσεις PyTorch, η πηγή προτείνει την επιλογή --no-build-isolation. Τα μοντέλα εκπαιδεύονται σε μεγάλα σύνολα δεδομένων, όπως το Pile και το SlimPajama, και καλύπτουν διαφορετικές υπολογιστικές ανάγκες.

Οι διεπαφές περιλαμβάνουν το επιλεκτικό επίπεδο SSM, το μπλοκ Mamba και πλήρεις δομές γλωσσικών μοντέλων. Το μπλοκ Mamba χρησιμοποιεί αιτιακό επίπεδο Conv1d και ενσωματώνεται σε νευρωνικά δίκτυα. Σημείωση για την πηγή: το αγγλικό άρθρο αναφέρεται σε παράδειγμα Python για δημιουργία μοντέλου και επεξεργασία δεδομένων, αλλά δεν περιλαμβάνει τον αντίστοιχο κώδικα.

Τα προεκπαιδευμένα μοντέλα στο Hugging Face κυμαίνονται από 130 εκατομμύρια έως 2,8 δισεκατομμύρια παραμέτρους. Έχουν εκπαιδευτεί στο Pile και στο SlimPajama, ακολουθούν τα πρότυπα διαστάσεων του GPT-3 και προορίζονται για διαφορετικές απαιτήσεις απόδοσης και υπολογιστικών πόρων. Η πηγή αναφέρει υψηλό ρυθμό επεξεργασίας και ακρίβεια σε εφαρμογές όπως η μοντελοποίηση γλώσσας.

Η επίδραση του Mamba

Το Mamba προσφέρει μια ισχυρή εναλλακτική στις αρχιτεκτονικές Transformer για δεδομένα πλούσια σε πληροφορία. Ο σχεδιασμός του αξιοποιεί σύγχρονο υλικό, βελτιστοποιώντας μνήμη και παράλληλη επεξεργασία. Ο δημόσιος κώδικας και τα προεκπαιδευμένα μοντέλα το καθιστούν προσιτό εργαλείο για ερευνητές και προγραμματιστές τεχνητής νοημοσύνης και βαθιάς μάθησης.