Googles forskare utvecklar ny AI-teknik som inte slösar hjärnkraft på värdelösa ord

Tabell över länkar

1. Introduktion
2. Bakgrund
3. Implementering av Mixture-of-Depths transformatorer

   • 3.1. Definiera en beräkningsbudget

   • 3.2. Ledning runt transformatorblock

   • 3.3. Rutningsscheman

   • 3.4. Implementering av routing

   • 3.5. Provtagning och 3.6. Träningsmetoder

4. Resultat
   • 4.1. Träning, isoFLOP jämförelser
   • 4.2. Autoregressiv utvärdering och 4.3. Blandning-av-djup-och-experter (MoDE)
5. Diskussion och referenser

Transformatorbaserade språkmodeller sprider FLOP:er enhetligt över inmatningssekvenser. I detta arbete visar vi att transformatorer istället kan lära sig att dynamiskt allokera FLOP:er (eller beräkna) till specifika positioner i en sekvens, och optimera allokeringen längs sekvensen för olika lager över modelldjupet. Vår metod upprätthåller en total beräkningsbudget genom att begränsa antalet tokens (𝑘) som kan delta i självuppmärksamhet och MLP-beräkningar vid ett givet lager. De tokens som ska bearbetas bestäms av nätverket med hjälp av en top-𝑘 routingmekanism. Eftersom 𝑘 definieras a priori, använder denna enkla procedur en statisk beräkningsgraf med kända tensorstorlekar, till skillnad från andra villkorliga beräkningstekniker. Ändå, eftersom identiteterna för 𝑘-symbolerna är flytande, kan denna metod förbruka FLOP:ar olikformigt över tids- och modelldjupdimensionerna. Således är beräkningsutgifter helt förutsägbara totalt sett, men dynamiska och kontextkänsliga på token-nivå. Inte bara modeller som tränas på detta sätt lär sig att dynamiskt allokera beräkningar, de gör det effektivt. Dessa modeller matchar baslinjeprestanda för likvärdiga FLOPS- och väggklockor för att träna, men kräver en bråkdel av FLOP:arna per framåtpassning, och kan vara uppåt 50 % snabbare att stega under provtagning efter träning.

1. Inledning

Alla problem kräver inte lika mycket tid eller ansträngning att lösa. Analogt, i språkmodellering kräver inte alla tokens och sekvenser samma tid eller ansträngning för att exakt göra en förutsägelse. Och ändå förbrukar transformatormodeller samma mängd beräkning per token i ett framåtpass. Idealiskt skulle transformatorer använda mindre totala beräkningsbudgetar genom att inte spendera beräkning i onödan.

Villkorsberäkning är en teknik som försöker minska den totala beräkningen genom att endast använda den när den behövs (Bengio et al., 2016; Bengio, 2013; Bengio et al., 2013). Olika algoritmer erbjuder lösningar på när och hur mycket beräkning som ska användas (Ainslie et al., 2023; Bapna et al., 2020; Fedus et al., 2022). Men generella formuleringar av detta utmanande problem kanske inte fungerar bra med befintliga hårdvarubegränsningar eftersom de tenderar att introducera dynamiska beräkningsgrafer (Dehghani et al., 2018; Graves, 2016). De mest lovande villkorliga beräkningsmetoderna kan istället vara de som är harmoniska med vår nuvarande hårdvarustack, som prioriterar statiska beräkningsgrafer, och kända tensorstorlekar som väljs för att maximera hårdvaruanvändningen.

Här överväger vi problemet med språkmodellering med hjälp av en statisk beräkningsbudget som kan göras mindre än den som används av en vaniljtransformator. Nätverket måste lära sig hur man dynamiskt allokerar tillgänglig beräkning genom att fatta beslut per token, i varje lager, om var man ska spendera beräkning från den tillgängliga budgeten. I vår implementering är total compute användardefinierad och oföränderlig före träning, snarare än att vara en funktion av nätverkets direkta beslut. Således kan maskinvarueffektivitetsvinster – som minskat minnesfotavtryck eller minskade FLOPs per framåtpassning – förutses och utnyttjas i förväg. Som vi kommer att visa kan dessa vinster uppnås utan att offra den totala prestandan.

Vi utnyttjar ett tillvägagångssätt som liknar Mixture of Experts (MoE) transformatorer, där dynamiska routingbeslut på tokennivå fattas över nätverksdjupet. Med utgångspunkt från MoE väljer vi att antingen tillämpa en beräkning på en token (som skulle vara fallet för en standardtransformator), eller att skicka den genom en restanslutning (förbli oförändrad och spara beräkning). Också i motsats till MoE, tillämpar vi denna routing på både framåt MLP och multi-head uppmärksamhet. Eftersom detta därför också påverkar de nycklar och frågor vi bearbetar, fattar routingen beslut inte bara om vilka tokens som ska uppdateras, utan också vilka tokens som görs tillgängliga att ta hand om. Vi hänvisar till denna strategi som Mixture-of-Depths (MoD) för att betona hur individuella tokens passerar genom olika antal lager, eller block, genom transformatorns djup (se figur 1).

MoD-tekniken gör det också möjligt att avväga prestanda med hastighet. Å ena sidan kan man träna en MoD-transformator som förbättrar vaniljtransformatorer med så mycket som 1,5 % på det slutliga log-sannolikhetsträningsmålet för likvärdiga tränings-FLOPs (isoFLOP), och samtidigt som det tar en motsvarande mängd väggklocka att träna. Å andra sidan kan man träna en MoD-transformator som uppnår träningsförlustparitet med en isoFLOP-optimal vaniljtransformator, men som använder en bråkdel av FLOP:arna (uppåt 50%) per framåtpassning, och därmed är snabbare att stega. Tillsammans innebär dessa resultat att MoD-transformatorer lär sig att dirigera intelligent (dvs. hoppa över beräkningar som är onödiga) eftersom de kan uppnå lika eller bättre loggsannolikheter per sekvens trots ett mindre FLOP-fotavtryck per framåtpassning.