NVIDIA rozwiązuje główny kompromis współczesnych modeli LLM

Pracując z dowolnym dużym modelem językowym (LLM), użytkownik musi wybierać między szybkością a jakością. Modele autoregresywne, takie jak GPT, zapewniają wysoką jakość, ale działają wolno, ponieważ generują po jednym tokenie na raz.

Modele dyfuzyjne działają odwrotnie. Generują całe bloki równolegle, co czyni je szybszymi, jednak ta prędkość odbywa się kosztem jakości.

To właśnie ten problem postanowili rozwiązać badacze z NVIDIA.

Wzięto model o rozmiarze 30B i podzielono go na dwie części, aby generować tokeny równolegle, zamiast sekwencyjnie. Tak powstał Nemotron-Labs-TwoTower – dyfuzyjny model językowy od NVIDIA Research, zaadaptowany na bazie Nemotron-3-Nano-30B-A3B.

Zanim jednak zrozumiesz, jak działa TwoTower, kluczowe jest uchwycenie przyczyn, dla których modele dyfuzyjne LLM ustępują miejsca modelom autoregresywnym.

Model dyfuzyjny LLM zaczyna od bloku zamaskowanych tokenów i odsłania je w kilku przejściach. W każdym przejściu czyta cały już wygenerowany tekst, jednocześnie przewidując zamaskowane tokeny w bieżącym bloku. Oznacza to, że ta sama sieć wykonuje dwa zupełnie różne zadania jednocześnie: musi rozumieć istniejący kontekst, przewidywać tokeny, które jeszcze nie istnieją.

Jeden zestaw wag jest skierowany do dwóch różnych zadań, więc sieć nie może w pełni wyspecjalizować się w żadnym z nich.

TwoTower rozwiązuje ten problem, nie zmuszając jednej sieci do wykonywania obu zadań. Zamiast tego, wstępnie wytrenowany model autoregresywny jest dzielony na dwie wieże (towers).

1. Wieża Kontekstu (Context Tower)

Jest to zamrożona część modelu, która czyta cały już wygenerowany tekst dokładnie tak, jak robił to oryginalny model. Nie jest ona dalej trenowana, dzięki czemu zdolności intelektualne modelu są w pełni zachowane.

2. Wieża Denoiser (Denoiser Tower)

Jest to trenowalna część modelu, która generuje każdy nowy blok tokenów równolegle. Aby utrzymać kontekst, stale wykorzystuje mechanizm cross-attention do Wieży Kontekstu.

Najciekawszą częścią rozwiązania jest sposób interakcji tych dwóch wież. Są one połączone warstwa po warstwie: piąta warstwa Wieży Denoiser współpracuje z piątą warstwą Wieży Kontekstu, szósta z szóstą i tak dalej. Pozwala to modułowi Denoiser uzyskać dostęp do pełnej, wielopoziomowej reprezentacji kontekstu z modelu bazowego, a nie tylko do jego końcowej ukrytej reprezentacji.

Wyniki:

• 2,42-krotnie wyższa przepustowość generowania
• Zachowanie 98,7% jakości oryginalnego modelu
• Architektura oparta na hybrydowym modelu Mamba–Transformer MoE o rozmiarze 30B
• Do treningu użyto jedynie około 2,1 biliona tokenów – stanowią one ułamek 25 bilionów tokenów wykorzystanych podczas pre-treningu oryginalnego modelu

Ta ostatnia liczba jest szczególnie istotna. Żaden z komponentów nie był trenowany od zera. Jest to adaptacja, którą można dodać do już istniejącego modelu autoregresywnego. Stosunkowo tanio przekształca powolne dekodowanie jednego tokenu na raz w szybkie, równoległe generowanie bloków.

Dlatego też w przypadku systemów wrażliwych na opóźnienia (latency-sensitive) jest to podejście, które zdecydowanie warto zbadać. Nie trzeba już wybierać między modelem, któremu ufamy, a prędkością, której potrzebujemy.

Wagi, kod oraz metodologia treningu są dostępne publicznie.

Artykuł naukowy:

https://arxiv.org/pdf/2606.26493