MiniMax’ inzet op lineaire aandacht: Een gesprek met Zhong Yiran, Hoofd MiniMax-01 Architectuur
Hoewel de Transformer-architectuur momenteel het generatieve AI-landschap domineert, zijn er beperkingen en komen er uitdagers op. MiniMax-01, met zijn gedurfde omarming van lineaire aandachtmechanismen en zijn uitbreiding tot een ongekende 456 miljard parameters, is zo’n disruptor in de open-source community. Dit is zowel een technologische gok als een potentiële mijlpaal in architecturale innovatie.
In dit interview spreken we met Zhong Yiran, het hoofd van de MiniMax-01 Architectuur, om de reis van lineaire aandacht van het laboratorium naar grootschalige industriële modellen te verkennen, evenals zijn gedachten en inzichten over modelarchitectuur.
Een Pionier op een Niet-Mainstream Technische Weg
Kunt u uzelf kort voorstellen?
Ik ben Zhong Yiran, Senior Research Director bij MiniMax, waar ik voornamelijk toezicht houd op het ontwerp van netwerkarchitecturen en grootschalige modellen voor multimodale begrip. Bij MiniMax is mijn belangrijkste verantwoordelijkheid het leiden van het ontwerp van de MiniMax-01 netwerkstructuur.
Voorheen was ik PI voor de New Architecture Exploration Group bij het Shanghai Artificial Intelligence Laboratory, waar ik me richtte op efficiënte trainingsmodelleringsmethoden voor niet-transformer architecturen en onderzoek naar visueel-audio-taal multimodale fusie.
Wanneer bent u begonnen met onderzoek naar lineaire aandacht, en waarom heeft u voor deze technische route gekozen?
Ik begon rond juli 2021 met onderzoek naar lineaire aandacht. Dit kwam voort uit een paper waaraan ik werkte voor mijn PhD in 2020, ‘Invertible Attention’. Destijds waren zowel inverseerbare neurale netwerken als aandachtmechanismen behoorlijk populair, dus we combineerden ze in ons onderzoek.
Later raakten sommige leden van ons team erg geïnteresseerd in wiskunde. Efficiënte sequentiemodelleringsmethoden zoals lineaire aandacht vereisen een sterke wiskundige basis en omvatten talrijke formule-afleidingen, wat perfect aansloot bij de interesses van het team, dus kozen we deze richting.
Wat was de status van lineaire aandacht in de industrie op dat moment?
Het was erg niet-mainstream, met weinig mensen die eraan werkten. De meeste onderzoekers waren gericht op transformers, die in wezen de dominante kracht in NLP waren geworden.
We dachten dat we, in plaats van zomaar een gezicht in de menigte te zijn die transformer-onderzoek deed, iets anders moesten doen.
Hoe heeft u het technische potentieel van de lineaire aandachtroute beoordeeld?
Onze initiële motivatie was eenvoudig: het aanpakken van de kwadratische computationele complexiteit van transformers. We hebben verschillende methoden getest, waaronder sparse transformers en lineaire aandacht.
We ontdekten dat sparse transformers wel werkten en een hogere snelheid en een lager geheugengebruik boden in vergelijking met transformers. Lineaire aandacht presteerde echter slecht en was ook traag. Desondanks kozen we ervoor om lineaire aandacht na te streven.
Een reden was de wiskundige aantrekkingskracht – we geloofden dat de prestaties beter zouden moeten zijn. De andere was dat we voelden dat de bovengrens van sparse aandacht volledige aandacht was, waardoor het moeilijk te overtreffen was. Lineaire aandacht daarentegen had het potentieel om het te overtreffen.
Kunt u uitleggen wat lineaire aandacht is?
Lineaire aandacht is in wezen een kerneltrick. In transformers omvat het vermenigvuldigen van de Q-, K- en V-matrices verschillende computationele complexiteiten, afhankelijk van of u QK eerst of KV eerst vermenigvuldigt, vanwege de verschillende dimensies.
Het eerst vermenigvuldigen van KV kan de computationele complexiteit tot lineair verminderen. Het probleem is echter dat QK-vermenigvuldiging wordt gevolgd door een softmax-bewerking, die niet voldoet aan de commutatieve eigenschap en niet gemakkelijk kan worden gesplitst in het eerst vermenigvuldigen van KV. Daarom is de eerste stap in lineaire aandacht het verwijderen van de softmax.
Maar het verwijderen van de softmax beïnvloedt de resultaten. De daaropvolgende taak is om de consistentie in de resultaten te behouden zonder softmax, en dat is wat lineaire aandacht beoogt te bereiken.
Wat zijn de fundamentele verschillen tussen lineaire aandacht, sparse aandacht en lineaire RNN-architecturen?
Sparse aandacht is in wezen nog steeds softmax-aandacht. Het berekent gewoon minder punten dan een dichte aandachtsmatrix. Sliding window attention berekent bijvoorbeeld alleen de aandachtsscore binnen een venster, waardoor versnelling wordt bereikt door de hoeveelheid berekeningen te verminderen.
Lineaire RNN’s en lineaire aandacht zijn in wezen hetzelfde, alleen noemen sommigen het RNN’s en anderen aandacht.
Alles kan in RNN-vorm worden geschreven. Lightning attention komt bijvoorbeeld overeen met RWKV-4, terwijl RWKV-7 een verbeterde versie is van het gated delta net. Hoewel ze in wezen hetzelfde zijn, verschillen hun implementatiedetails.
Wat zijn de belangrijkste mijlpalen in het onderzoek naar lineaire aandachtmechanismen?
Rond 2018-19 toonde onderzoek aan dat de theoretische computationele complexiteit van transformer softmax-aandacht kon worden verminderd met behulp van kerneltricks, maar de resultaten waren slecht en de efficiëntie was laag.
In 2019-20 was sparse aandacht dominant, met bedrijven als Google die veel sparse aandachtvarianten voorstelden. Later begon lineaire aandacht op te komen, maar het stond voor de uitdaging van slechte prestaties en lage snelheid.
Onderzoekers hanteerden voornamelijk twee benaderingen voor verbetering: de ene was het benaderen van de softmax-functie, waardoor de distributie overeenkomt met softmax; de andere, die wij kozen, was het modelleren met volledig verschillende methoden, zonder ons zorgen te maken over het benaderen van softmax.
We publiceerden ons eerste paper, ‘COSFORMER: RETHINKING SOFTMAX IN ATTENTION’, in oktober 2021, dat de softmax-bewerking verving door een cosinusfunctie, waardoor de berekening kon worden gesplitst.
In de eerste helft van 2022 publiceerden we een tweede paper, ‘The Devil in Linear Transformer’, dat de redenen voor de prestatievermindering van lineaire aandacht analyseerde en oplossingen bood. Dit was de voorloper van lightning attention.
Later deden we ook onderzoek naar positie-encodings specifiek voor lineaire aandacht en lange convoluties, en publiceerden we TNN, ‘TOEPLITZ NEURAL NETWORK FOR SEQUENCE MODELING’, een methode die vergelijkbaar is met S4 (de voorganger van Mamba).
Ten slotte lanceerden we lightning attention, dat de prestaties van transformers evenaarde door middel van verbeterde vervalmethoden en netwerkstructuren. We gebruikten ook een tiling-techniek om het sneller te maken.
Wat zijn uw gedachten over de huidige niet-transformer architectuur technische routes?
Lineaire aandacht is eigenlijk een niet-transformer methode. Momenteel zijn, naast RNN-achtige benaderingen, andere niet-transformer architecturen in verval.
CNN’s zoals lange convoluties en grote kernelconvoluties voelen bijvoorbeeld alsof ze geleidelijk zijn geëlimineerd vanwege slechte prestaties, maar ze zijn eigenlijk behoorlijk sterk in bepaalde aspecten en hebben nog steeds enig effect in sequentiemodellering, zoals anomaliedetectietaken.
Er zijn eigenlijk maar drie niet-transformer architecturen: lineaire aandacht, lange convoluties en lineaire RNN’s.
Maar in werkelijkheid kunnen deze drie worden verenigd tot één, die we het lineaire complexiteitsmodel noemen. We schreven een artikel dat alle drie omvat.
Wat zijn de belangrijkste verschillen tussen lightning attention en Mamba en RWKV?
Het meest kernverschil is dat lightning attention de eenvoudigste lineaire aandacht is. Mamba en RWKV gebruiken beide data-afhankelijk verval, terwijl lightning attention handgemaakt verval gebruikt voor snelheid.
Hoewel leerbaar verval betere resultaten kan opleveren, offert het snelheid op. RWKV-7 is bijvoorbeeld 10-15% langzamer dan het gating delta net, terwijl het gated delta net ongeveer de helft van de snelheid van lightning attention heeft.
Het modelleringseffect van RWKV is inderdaad beter dan lightning attention, maar het is langzamer en heeft het retrieval-probleem nog niet opgelost.
Is het nu een industrieel consensus dat lineaire aandacht een hoge en haalbare bovengrens heeft?
Nee, als het consensus zou zijn, zou iedereen lineaire aandachtmodellen opschalen. En het is nu ook geen consensus. Als het zo was, zou iedereen lineair doen, maar zoals je kunt zien, is dat niet het geval.
Maar voor ons zagen we dit al in de tweede helft van 2023. Destijds vroeg ik veel mensen en praatte ik met velen, en het meest voorkomende punt dat ze aanhaalden was dat ze wisten dat lineaire aandacht op kleine schaal wel werkte, maar ze voelden dat het zou falen zodra het werd opgeschaald.
Destijds dacht ik dat ik het zou opschalen zodat iedereen het kon zien. Nu MiniMax-01 uit is, twijfelt niemand aan het vermogen van lineaire aandacht op grote schaal.
Van Kleine Experimenten tot Grootschalige Implementatie
Denkt u dat de bovengrens van lineaire aandacht volledige aandacht kan overtreffen?
We kunnen nu zien dat hybride architecturen beter zijn dan pure transformers. Maar het grootste probleem met pure lineaire aandacht is retrieval-vermogen, wat een moeilijk probleem is voor de academische wereld om op te lossen.
Bestaande methoden, hoewel complex en traag, kunnen het nog steeds niet volledig oplossen, en daarom is het noodzakelijk om over te stappen op hybride architecturen.
Welk knooppunt heeft u waargenomen waardoor u besloot om uit het lab te komen?
In mei-juni 2023 hadden we intern al lightning attention 2, wat de eerste lineaire aandachtimplementatie ter wereld was die sneller was dan Flash attention.
Wij zijn van mening dat het de industriële rode lijn heeft overschreden, en de technologische volwassenheid is zeer hoog en kan worden opgeschaald.
Hoe definieert u deze industriële rode lijn?
Ten eerste is het effect beter dan transformer, en ten tweede is het sneller dan transformer. Dit geeft het de mogelijkheid om transformer te vervangen. We hebben dit destijds geverifieerd op een 15B schaal dicht model.
Op het knooppunt toen u uit het lab kwam, waarom bent u uiteindelijk samengekomen met MiniMax?
Ik had destijds eigenlijk met enkele grote bedrijven gesproken. Maar uiteindelijk heb ik dit toch met MiniMax laten gebeuren.
Allereerst is cosformer een artikel waar ik met Junjie aan heb samengewerkt. We hebben een basis voor samenwerking. Junjie was mijn baas toen hij bij SenseTime was. Aan het einde van 23 nodigde Junjie me uit voor een diner. Hij heeft meer vertrouwen in de mogelijkheden van deze geavanceerde technologieën. Mijn inzicht is dat hij destijds ook op zoek was naar een technische doorbraak.
Destijds had MiniMax het onderzoek naar Moe voltooid, en er waren eigenlijk maar heel weinig technische doorbraakpunten voor de volgende stap. Destijds was lightning attention uitgebracht en was mamba ook populair, dus in zijn ogen was het een haalbare richting.
Heeft dit te maken met het interactieve companion product van MiniMax?
Er is geen verband. Yan Junjie maakt zich meer zorgen over debovengrens van het model en hoe deze grens verder te doorbreken.
Lineaire aandacht is in de publieke opinie misschien meer een richting om de efficiëntie te doorbreken, in plaats van de bovengrens te doorbreken.
Het punt hier is dat ten eerste de rekenkracht van elke fabrikant constant is. Hoe sneller het model kan worden versneld, hoe meer gegevens het kan eten, en hoe beter het geproduceerde model. Wanneer de rekenkracht constant is, hoe sneller het model, hoe beter.
Heeft u een situatie waargenomen waarin gegevens een piek hebben bereikt?
Nog niet, toch? Data bevindt zich nog steeds in de fase van continue opschaling, maar het is misschien niet zo agressief als in 23.
Omdat de data altijd toeneemt en er elke dag nieuwe data uitkomt. Voor het model zijn er elke dag nieuwe data te verwerken. De data die elke dag door internet wordt geproduceerd is zoveel. Door middel van opschoning kunnen we er nog steeds nieuwe data uit halen.
Is de datagroei afgenomen vergeleken met de data die al zoveel jaren van menselijke ontwikkeling bestaat?
Eigenlijk niet noodzakelijk. Kijk naar de vijfduizend jaar van de Chinese geschiedenis, en alleen die paar boeken zijn verzameld. Maar met de ontwikkeling van internet is de toename van het datavolume een zeer steile curve. De totale data die vóór internet is gegenereerd is misschien niet zo veel als de data die in één jaar later is gegenereerd.
Met welke uitdagingen werd lightning attention geconfronteerd tijdens het opschalen?
Om de schaalbaarheid te verifiëren, hebben we eerst schaalwetexperimenten gedaan, waarbij we geleidelijk uitbreidden van kleine modellen tot 7B, 9B, en uiteindelijk opschaalden naar modellen met meer dan 400B.
En we hebben theoretisch bewezen dat de capaciteit van lineair groter is dan die van transformer.
We definiëren capaciteit als de grootte van de huidige staten van de RNN. Voor transformer is de capaciteitsgrootte O(d), waarbij d de grootte is; voor lineaire aandacht is de capaciteitsgrootte d²/h. Aangezien d veel groter is dan h, is de capaciteit groter.
Uiteindelijk hebben we ook geverifieerd dat het hybride model beter is dan de pure transformer.
Hoe wordt de 4M lengte sequentie venster bereikt?
Voor lightning kan de trainingslengte willekeurig zijn. Zolang de rekenkracht volledig wordt benut, is de snelheid van het trainen van 8K, 32K of 128K hetzelfde, en is de TGS (token per GPU per seconde) hetzelfde.
Omdat transformer n² computationele complexiteit is, hoe langer de sequentie, hoe sneller de computationele complexiteit groeit, en de latentie neemt toe in een kwadratische curve. Op 1M lengte is de latentie van softmax aandacht 2.700 keer zo hoog als die van lightning attention.
Welke technische uitdagingen moeten nog worden aangepakt om in de toekomst een oneindig contextvenster te bereiken?
In onze huidige hybride architectuur is er nog steeds 1/8 softmax aandacht. Dit is een knelpunt op 1M lengte. De latentie die door deze 1/8 wordt veroorzaakt is veel hoger dan de overige 7/8 van lineaire aandacht.
Als we lange teksten willen optimaliseren, moeten we overwegen om het softmax-aandachtsgedeelte te optimaliseren. We kunnen leren van sparse aandachtmethoden om het sneller en lichter te maken.
Daarnaast overwegen we ook om de mengverhouding van softmax en lineaire aandacht extremer te maken, niet langer 1/8, maar mogelijk 1/16 of 1/32. De meest radicale oplossing is om slechts één laag softmax in het hele model te plaatsen, maar ter verzekering hebben we dit niet aangenomen, voornamelijk rekening houdend met de impact op het retrieval-vermogen.
Waarom is retrieval-vermogen zo belangrijk voor het model?
Retrieval is de basis van in-context learning en is een noodzakelijke voorwaarde.
Je moet de informatie in de context onthouden om in-context learning te doen, en in-context learning is de basis van alle geavanceerde mogelijkheden van de huidige grote modellen, zoals CoT (Chain of Thought), vooral lange CoT, die allemaal afhankelijk zijn van retrieval-vermogen.
Doorslaggevende Nieuwe Architectuur
Heeft u aandacht besteed aan de nieuwste architecturale verbeteringen in FFN en aandacht in de industrie?
De verbetering van FFN is Moe. Ik heb ook aandacht besteed aan Byte’s Ultra Mem, maar ik denk dat het een lossy ding is, een lossy compressie. Er kunnen problemen zijn als het in de toekomst wordt opgeschaald, maar we hebben het niet opgeschaald, dus ik kan alleen zeggen dat er problemen kunnen zijn.
Omdat FFN in principe deze zijn. Onze verbeteringen in het Moe-gebied zijn niets meer dan het veranderen van de vorige grote expert naar de huidige kleine expertmodus, waardoor het schaarser wordt, en vervolgens wat versnelling doen, wat verder onderzoek vereist.
Als je het verder wilt optimaliseren, omdat FFN matrixvermenigvuldiging is, kan de optimalisatie alleen op het CUDA-niveau door Nvidia worden gedaan, door enkele van de onderliggende optimalisaties van matrixvermenigvuldiging te doen.
Heeft u aandacht besteed aan de verbeteringen in de aandachtsarchitectuur in de industrie?
De verbeteringen aan aandacht zijn in principe lineair. We overwegen ook of we in de toekomst een sterkere Linear willen maken, en lineaire aandacht verder versnellen op de huidige basis.
Er zijn veel manieren om te verbeteren, de ene is om het verval te veranderen, en de andere is om enkele kleine trucjes binnenin te veranderen. U kunt uitkijken naar ons nieuwe paper.
Is onze huidige verhouding tussen contextlengte en inferentiekosten relatief geavanceerd?
Zodra het gaat om het verlengen van de sequentielengte, hebben we een zeer duidelijk voordeel in rekenkrachtkosten. Hoe langer het is, hoe duidelijker het kostenvoordeel zal zijn, of het nu inferentie of training is.
Op 1M bijvoorbeeld is de rekenkracht die door lineaire aandacht wordt verbruikt 1/2700 van volledige aandacht. Ter vergelijking, omdat we nog steeds 1/8 van volledige aandacht hebben, is het in principe 1/8 van de transformer-architectuur, omdat lineaire aandacht in principe niet als een uitgave telt.
Als de berekeningskosten zo laag zijn, kan het dan een berekeningsknelpunt bereiken?
Nu is het inderdaad een knelpunt voor geheugentoegang. Decoderen is een knelpunt voor geheugentoegang, geen berekeningsknelpunt. Omdat lightning erg snel is, is het te snel om geheugentoegang toe te staan om zo weinig resources te gebruiken als berekening. Dit komt vooral omdat de sequentielengte in daadwerkelijke toepassingen niet lang genoeg is.
Hoe je er in de toekomst een berekeningsknelpunt van kunt maken, hangt af van hoe je geheugentoegang kunt optimaliseren. Dit zullen dingen zijn waar de technische afdeling verantwoordelijk voor moet zijn.
Als lineaire architectuur de mainstream architectuur van de volgende generatie wordt, welke hardware aanpassingsverbeteringen zouden er meer geschikt voor zijn?
Een zeer lastig ding hier is dat we rekening moeten houden met de sequentielengte. Als uw sequentielengte gericht is op 8K of 32K, dan is aandacht slechts iets meer dan tien procent, en de overige tachtig procent is het FFN-gedeelte.
Zelfs als je aandacht tot het uiterste optimaliseert, tot 0, heb je slechts iets meer dan tien procent van de latentie geoptimaliseerd. Maar als je de sequentielengte verlengt, zal het aandeel van aandacht groter en groter worden. Dit is vergeleken met volledige aandacht, maar voor lineaire aandacht blijft het aandeel ongewijzigd.
Omdat FFN ook lineair is, en lineaire aandacht ook lineair is, is het aandeel ongeveer 10%, wat bijna ongewijzigd is, zelfs in het geval van 1M.
Maar als het volledige aandacht is, kan de aandachts