Výzkumníci Google vyvíjejí novou technologii umělé inteligence, která neplýtvá mozkem na zbytečná slova

Tabulka odkazů

1. Zavedení
2. Pozadí
3. Implementace smíšených hloubkových transformátorů

   • 3.1. Definování výpočetního rozpočtu

   • 3.2. Vedení kolem transformátorových bloků

   • 3.3. Schémata směrování

   • 3.4. Implementace směrování

   • 3.5. Odběr vzorků a 3.6. Tréninkové metody

4. Výsledky
   • 4.1. Trénink, srovnání isoFLOP
   • 4.2. Autoregresivní hodnocení a 4.3. Směs hloubek a odborníků (MoDE)
5. Diskuse a reference

Jazykové modely založené na transformátoru rozdělují FLOPy rovnoměrně napříč vstupními sekvencemi. V této práci demonstrujeme, že transformátory se místo toho mohou naučit dynamicky přidělovat FLOP (nebo počítat) do konkrétních pozic v sekvenci a optimalizovat alokaci podél sekvence pro různé vrstvy v hloubce modelu. Naše metoda vynucuje celkový výpočetní rozpočet omezením počtu tokenů (𝑘), které se mohou podílet na výpočtech sebepozornosti a MLP na dané vrstvě. Tokeny, které mají být zpracovány, určuje síť pomocí top-𝑘 směrovacího mechanismu. Protože 𝑘 je definováno a priori, používá tento jednoduchý postup na rozdíl od jiných technik podmíněného výpočtu statický výpočtový graf se známými velikostmi tenzorů. Nicméně, protože identity 𝑘 tokenů jsou proměnlivé, může tato metoda utratit FLOPy nerovnoměrně napříč časem a dimenzemi hloubky modelu. Výpočetní výdaje jsou tedy v součtu zcela předvídatelné, ale na úrovni tokenů jsou dynamické a závislé na kontextu. Takto trénované modely se nejen naučí dynamicky alokovat výpočet, ale dělají to efektivně. Tyto modely odpovídají základnímu výkonu pro ekvivalentní FLOPS a nástěnné hodiny pro trénování, ale vyžadují zlomek FLOPů na dopředný průchod a mohou být až o 50 % rychlejší v kroku během vzorkování po tréninku.

1. Úvod

Ne všechny problémy vyžadují stejné množství času nebo úsilí k vyřešení. Analogicky, v jazykovém modelování ne všechny tokeny a sekvence vyžadují stejný čas nebo úsilí k přesnému provedení předpovědi. A přesto modely transformátorů vynaloží stejné množství výpočtů na token v dopředném průchodu. V ideálním případě by transformátory využívaly menší celkové výpočetní rozpočty tím, že by zbytečně nevynakládaly výpočetní prostředky.

Podmíněný výpočet je technika, která se snaží snížit celkový výpočet tím, že jej vynakládá pouze v případě potřeby (Bengio et al., 2016; Bengio, 2013; Bengio et al., 2013). Různé algoritmy nabízejí řešení, kdy a kolik výpočetních prostředků by se mělo používat (Ainslie et al., 2023; Bapna et al., 2020; Fedus et al., 2022). Obecné formulace tohoto náročného problému však nemusí dobře fungovat se stávajícími hardwarovými omezeními, protože mají tendenci zavádět dynamické výpočetní grafy (Dehghani et al., 2018; Graves, 2016). Nejslibnější podmíněné výpočetní metody mohou být místo toho ty, které jsou v souladu s naším současným hardwarovým zásobníkem, který upřednostňuje statické výpočetní grafy a známé velikosti tenzorů, které jsou vybrány tak, aby maximalizovaly využití hardwaru.

Zde uvažujeme o problému jazykového modelování pomocí statického výpočetního rozpočtu, který může být menší než ten, který používá vanilkový transformátor. Síť se musí naučit, jak dynamicky přidělovat dostupné výpočty tím, že se pro každý token v každé vrstvě rozhoduje o tom, kam utratit výpočet z dostupného rozpočtu. V naší implementaci je celkový výpočet uživatelem definovaný a neměnný před školením, spíše než aby byl funkcí síťových rozhodnutí za chodu. Zvýšení efektivity hardwaru – jako je snížená paměťová náročnost nebo snížené FLOP na jeden dopředný průchod – lze tedy očekávat a využít s předstihem. Jak si ukážeme, těchto zisků lze dosáhnout bez obětování celkového výkonu.

Využíváme přístup podobný transformátorům Mixture of Experts (MoE), ve kterých se rozhodnutí o dynamickém směrování na úrovni tokenů provádějí v celé hloubce sítě. Vycházíme-li z MŽP, volíme buď použít výpočet na token (jako by tomu bylo v případě standardního transformátoru), nebo jej předat přes zbytkové spojení (zůstat beze změny a šetřit výpočet). Také na rozdíl od MŽP toto směrování aplikujeme jak na dopředné MLP, tak na vícehlavou pozornost. Protože to má dopad také na klíče a dotazy, které zpracováváme, směrování rozhoduje nejen o tom, které tokeny aktualizovat, ale také o tom, které tokeny jsou zpřístupněny pro zpracování. Tuto strategii nazýváme Mixture-of-Depths (MoD), abychom zdůraznili, jak jednotlivé tokeny procházejí různým počtem vrstev nebo bloků hloubkou transformátoru (viz obrázek 1).

Technika MoD také umožňuje vyrovnat výkon s rychlostí. Na jedné straně lze trénovat transformátor MO, který se oproti vanilkovým transformátorům zlepšuje až o 1,5 % v konečném cíli trénování logaritmické pravděpodobnosti pro ekvivalentní tréninkové FLOPy (isoFLOP), a přitom trénování zabere ekvivalentní množství času nástěnných hodin. Na druhou stranu lze trénovat MOD transformátor, který dosahuje parity ztrát při tréninku s isoFLOP optimálním vanilkovým transformátorem, ale který využívá zlomek FLOPů (až 50 %) na dopředný průchod, a proto je rychlejší v kroku. Tyto výsledky společně naznačují, že transformátory MO se učí inteligentně směrovat (tj. přeskakovat výpočty, které jsou zbytečné), protože mohou dosáhnout stejných nebo lepších log pravděpodobnosti na sekvenci i přes menší stopu FLOP na dopředný průchod.