Ascensiunea Mixture-of-Experts pentru Modele de Limbaj Mari și Eficiente

În lumea prelucrării limbajului natural (NLP), urmărirea construirii unor modele de limbaj mai mari și mai capabile a fost o forță motrice în spatele multor progrese recente. Cu toate acestea, pe măsură ce aceste modele cresc în dimensiune, cerințele computaționale pentru antrenare și inferență devin din ce în ce mai solicitante, împingând limitele resurselor hardware disponibile.

Apare Mixture-of-Experts (MoE), o tehnică care promite să alineze această povară computațională, permițând antrenarea unor modele de limbaj mai mari și mai puternice. Mai jos, vom discuta despre MoE, vom explora originile, mecanismele interne și aplicațiile sale în modelele de limbaj bazate pe transformatori.

Originile Mixture-of-Experts

Conceptul de Mixture-of-Experts (MoE) poate fi urmărit până în anii ’90, când cercetătorii au explorat ideea de calcul condițional, în care părți ale unei rețele neuronale sunt activate selectiv pe baza datelor de intrare. Una dintre lucrările pionieratice în acest domeniu a fost articolul “Adaptive Mixture of Local Experts” de Jacobs et al. din 1991, care a propus un cadru de învățare supervizată pentru un ansamblu de rețele neuronale, fiecare specializată într-o regiune diferită a spațiului de intrare.

Ideea de bază din spatele MoE este de a avea multiple rețele “expert”, fiecare responsabilă pentru prelucrarea unei submulțimi a datelor de intrare. Un mecanism de porționare, de obicei o rețea neuronală în sine, determină care expert(i) ar trebui să prelucreze o anumită intrare. Acest abordaj permite modelului să aloce resursele sale computaționale mai eficient, activând doar experții relevanți pentru fiecare intrare, în loc de a utiliza întreaga capacitate a modelului pentru fiecare intrare.

De-a lungul anilor, diverși cercetători au explorat și extins ideea de calcul condițional, conducând la dezvoltări precum MoE ierarhice, aproximații de rang scăzut pentru calculul condițional și tehnici pentru estimarea gradientelor prin neuroni stohastici și funcții de activare cu prag dur.

Mixture-of-Experts în Transformatori

Mixture of Experts

Deși ideea de MoE a existat de decenii, aplicarea sa la modele de limbaj bazate pe transformatori este relativ recentă. Transformatorii, care au devenit standardul de facto pentru modele de limbaj de ultimă generație, sunt compuși din multiple straturi, fiecare conținând un mecanism de atenție auto și o rețea neuronală feed-forward (FFN).

Inovația cheie în aplicarea MoE la transformatori constă în înlocuirea straturilor FFN dense cu straturi MoE sparse, fiecare compus din multiple rețele expert FFN și un mecanism de porționare. Mecanismul de porționare determină care expert(i) ar trebui să prelucreze fiecare token de intrare, permițând modelului să activeze selectiv doar o submulțime de experți pentru o anumită secvență de intrare.

Una dintre lucrările timpurii care a demonstrat potențialul MoE în transformatori a fost articolul “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” de Shazeer et al. din 2017. Această lucrare a introdus conceptul de strat MoE cu poartă rară, care a utilizat un mecanism de porționare care a adăugat raritate și zgomot la procesul de selecție a experților, asigurând că doar o submulțime de experți este activată pentru fiecare intrare.

De atunci, numeroase alte lucrări au avansat și mai mult aplicarea MoE la transformatori, abordând provocări precum instabilitatea antrenării, echilibrarea încărcării și inferența eficientă. Exemple notabile includ Switch Transformer (Fedus et al., 2021), ST-MoE (Zoph et al., 2022) și GLaM (Du et al., 2022).

Beneficiile Mixture-of-Experts pentru Modele de Limbaj

Beneficiul principal al utilizării MoE în modele de limbaj este capacitatea de a mări dimensiunea modelului, menținând în același timp un cost computațional relativ constant în timpul inferenței. Prin activarea selectivă a doar unei submulțimi de experți pentru fiecare token de intrare, modelele MoE pot atinge puterea de exprimare a unor modele dense mult mai mari, necesitând semnificativ mai puțină computație.

De exemplu, considerați un model de limbaj cu un strat FFN dens de 7 miliarde de parametri. Dacă înlocuim acest strat cu un strat MoE compus din opt experți, fiecare cu 7 miliarde de parametri, numărul total de parametri crește la 56 de miliarde. Cu toate acestea, în timpul inferenței, dacă activăm doar doi experți pe token, costul computațional este echivalent cu al unui model dens de 14 miliarde de parametri, deoarece se efectuează două multiplicări de matrice de 7 miliarde de parametri.

Această eficiență computațională în timpul inferenței este deosebit de valoroasă în scenariile de implementare în care resursele sunt limitate, cum ar fi dispozitivele mobile sau mediile de calcul edge. În plus, cerințele computaționale reduse în timpul antrenării pot conduce la economii substanțiale de energie și o amprentă de carbon mai mică, aliniindu-se cu accentul crescând pe practicile durabile de AI.

Provocări și Considerații

Deși modelele MoE oferă beneficii convingătoare, adoptarea și implementarea lor vin și cu mai multe provocări și considerații:

1. Instabilitatea Antrenării: Modelele MoE sunt cunoscute pentru a fi mai predispuse la instabilități de antrenare în comparație cu omologii lor denși. Această problemă apare din cauza naturii sparse și condiționale a activărilor experților, care poate duce la provocări în propagarea gradientului și convergență. Tehnici precum router z-loss (Zoph et al., 2022) au fost propuse pentru a mitigă aceste instabilități, dar este nevoie de cercetări suplimentare.
2. Finetuning și Supraantrenare: Modelele MoE tind să supraantreneze mai ușor în timpul finetuning-ului, especially când sarcina downstream are un set de date relativ mic. Acest comportament este atribuit capacității crescute și rarității modelului MoE, care poate duce la o specializare excesivă pe datele de antrenare. Strategii de regularizare și finetuning atente sunt necesare pentru a mitigă această problemă.
3. Cerințe de Memorie: Deși modelele MoE pot reduce costurile computaționale în timpul inferenței, ele au adesea cerințe de memorie mai mari comparativ cu modelele dense de dimensiuni similare. Acest lucru se datorează faptului că toate greutățile experților trebuie încărcate în memorie, chiar dacă doar o submulțime este activată pentru fiecare intrare. Constrângerile de memorie pot limita scalabilitatea modelului MoE pe dispozitive cu resurse limitate.
4. Echilibrarea Încărcării: Pentru a atinge eficiența computațională optimă, este crucial să se echilibreze încărcarea între experți, asigurându-se că niciun expert nu este supraîncărcat, în timp ce alții rămân subutilizați. Acest echilibru al încărcării se realizează de obicei prin pierderi auxiliare în timpul antrenării și prin reglarea atentă a factorului de capacitate, care determină numărul maxim de tokeni care pot fi asignați fiecărui expert.
5. Supraîncărcarea Comunicației: În scenariile de antrenare și inferență distribuite, modelele MoE pot introduce o supraîncărcare suplimentară de comunicare din cauza necesității de a schimba informații de activare și gradient între experți rezidenți pe dispozitive sau acceleratoare diferite. Strategii de comunicare eficiente și proiectarea modelului conștientă de hardware sunt esențiale pentru a mitigă această supraîncărcare.

În ciuda acestor provocări, beneficiile potențiale ale modelului MoE în permiterea unor modele de limbaj mai mari și mai capabile au stimulat eforturi semnificative de cercetare pentru a aborda și a mitigă aceste probleme.

Exemplu: Mixtral 8x7B și GLaM

Pentru a ilustra aplicarea practică a MoE în modele de limbaj, să considerăm două exemple notabile: Mixtral 8x7B și GLaM.

Mixtral 8x7B este o variantă MoE a modelului de limbaj Mistral, dezvoltat de Anthropic. Acesta constă din opt experți, fiecare cu 7 miliarde de parametri, rezultând un total de 56 de miliarde de parametri. Cu toate acestea, în timpul inferenței, se activează doar doi experți pe token, reducând efectiv costul computațional la cel al unui model dens de 14 miliarde de parametri.

Mixtral 8x7B a demonstrat performanțe impresionante, depășind modelul Llama de 70 de miliarde de parametri, oferind în același timp timp de inferență mult mai rapid. O versiune instructată a lui Mixtral 8x7B, numită Mixtral-8x7B-Instruct-v0.1, a fost lansată și pentru a-și îmbunătăți capacitățile de a urma instrucțiuni de limbaj natural.

Un alt exemplu remarcabil este GLaM (Google Language Model), un model MoE de largă scală dezvoltat de Google. GLaM utilizează o arhitectură de transformator doar cu decodificator și a fost antrenat pe un set masiv de date de 1,6 trilioane de tokeni. Modelul atinge performanțe impresionante la evaluările few-shot și one-shot, egalând calitatea GPT-3, dar utilizând doar o treime din energia necesară pentru antrenarea GPT-3.

Succesul lui GLaM se datorează în mare măsură arhitecturii sale MoE eficiente, care a permis antrenarea unui model cu un număr vast de parametri, menținând în același timp cerințe computaționale rezonabile. Modelul a demonstrat, de asemenea, potențialul modelului MoE de a fi mai eficient din punct de vedere energetic și mai durabil din punct de vedere al mediului, comparativ cu omologii săi denși.

Arhitectura Grok-1

GROK MIXTURE OF EXPERT

Grok-1 este un model MoE bazat pe transformator, cu o arhitectură unică proiectată pentru a maximiza eficiența și performanța. Să explorăm specificațiile cheie:

1. Parametri: Cu un număr impresionant de 314 miliarde de parametri, Grok-1 este cel mai mare model LLM deschis până în prezent. Cu toate acestea, datorită arhitecturii MoE, doar 25% din greutăți (aproximativ 86 de miliarde de parametri) sunt active în orice moment, îmbunătățind capacitățile de procesare.
2. Arhitectură: Grok-1 utilizează o arhitectură Mixture-of-8-Experts, cu fiecare token prelucrat de doi experți în timpul inferenței.
3. Straturi: Modelul constă din 64 de straturi de transformator, fiecare incorporând atenție multi-cap și blocuri dense.
4. Tokenizare: Grok-1 utilizează un tokenizator SentencePiece cu o dimensiune a vocabularului de 131.072 de tokeni.
5. Încorporări și Codificări Poziționale: Modelul prezintă încorporări de 6.144 de dimensiuni și utilizează încorporări poziționale rotative, permițând o interpretare mai dinamică a datelor în comparație cu încorporările poziționale fixe tradiționale.
6. Atenție: Grok-1 utilizează 48 de capete de atenție pentru interogări și 8 capete de atenție pentru chei și valori, fiecare cu o dimensiune de 128.
7. Lungimea Contextului: Modelul poate prelucra secvențe de până la 8.192 de tokeni în lungime, utilizând precizia bfloat16 pentru computație eficientă.

Performanță și Detalii de Implementare

Grok-1 a demonstrat performanțe impresionante, depășind LLaMa 2 70B și Mixtral 8x7B cu un scor MMLU de 73%, demonstrând eficiența și acuratețea sa pe diverse teste.

Cu toate acestea, este important de remarcat că Grok-1 necesită resurse GPU semnificative datorită dimensiunii sale impresionante. Implementarea curentă în lansarea sursă deschisă se axează pe validarea corectitudinii modelului și utilizează o implementare ineficientă a stratului MoE pentru a evita nevoia de kerneluri personalizate.

Totuși, modelul suportă fragmentarea activării și cuantificarea pe 8 biți, care pot optimiza performanța și reduce cerințele de memorie.

Într-o mișcare remarcabilă, xAI a lansat Grok-1 sub licența Apache 2.0, făcând greutățile și arhitectura sa accesibile comunității globale pentru utilizare și contribuții.

Lansarea sursă deschisă include un depozit de cod exemple JAX care demonstrează cum să încărcați și să rulați modelul Grok-1. Utilizatorii pot descărca greutățile punctului de control utilizând un client de torrent sau direct prin HuggingFace Hub, facilitând accesul ușor la acest model de ultimă generație.

Viitorul Mixture-of-Experts în Modele de Limbaj

Pe măsură ce cererea pentru modele de limbaj mai mari și mai capabile continuă să crească, adoptarea tehnicilor MoE este de așteptat să câștige și mai multă impuls. Eforturile de cercetare în curs se axează pe abordarea provocărilor rămase, cum ar fi îmbunătățirea stabilității antrenării, mitigarea supraantrenării în timpul finetuning-ului și optimizarea cerințelor de memorie și comunicare.

O direcție promițătoare este explorarea arhitecturilor MoE ierarhice, în care fiecare expert este la rândul său compus din mai mulți sub-experți. Acest abordaj ar putea permite o scalabilitate și o eficiență computațională și mai mare, menținând puterea de exprimare a modelelor mari.

În plus, dezvoltarea sistemelor de hardware și software optimizate pentru modele MoE este o zonă activă de cercetare. Acceleratoare specializate și cadre de antrenare distribuite proiectate pentru a gestiona eficient modelele de calcul sparse și condiționale ale modelelor MoE ar putea îmbunătăți și mai mult performanța și scalabilitatea acestora.

Mai mult, integrarea tehnicilor MoE cu alte progrese în modelarea limbajului, cum ar fi mecanismele de atenție sparse, strategiile eficiente de tokenizare și reprezentările multi-modale, ar putea duce la modele de limbaj și mai puternice și versatile, capabile să abordeze o gamă largă de sarcini.

Concluzie

Tehnica Mixture-of-Experts a apărut ca un instrument puternic în căutarea unor modele de limbaj mai mari și mai capabile. Prin activarea selectivă a experților pe baza datelor de intrare, modelele MoE oferă o soluție promițătoare la provocările computaționale asociate cu mărirea dimensiunii modelelor dense. Deși există încă provocări de depășit, cum ar fi instabilitatea antrenării, supraantrenarea și cerințele de memorie, beneficiile potențiale ale modelului MoE în ceea ce privește eficiența computațională, scalabilitatea și durabilitatea mediului îl fac o zonă de cercetare și dezvoltare interesantă.

Pe măsură ce domeniul prelucrării limbajului natural continuă să împingă limitele a ceea ce este posibil, adoptarea tehnicilor MoE este probabil să joace un rol crucial în permiterea următoarei generații de modele de limbaj. Prin combinarea MoE cu alte progrese în arhitectura modelului, tehnici de antrenare și optimizarea hardware-ului, putem aștepta modele de limbaj și mai puternice și versatile, capabile să înțeleagă și să comunice cu oamenii într-un mod natural și fără efort.