Neues Werk des Teams von Han Song bei NVIDIA: Effizientes Sprachmodell mit nachträglichem Neural Architecture Search

NVIDIA hat wieder ein sensationelles Open-Source-Angebot herausgebracht!

Das Team von Han Song hat ein brandneues, effizientes Sprachmodell namens Jet-Nemotron basierend auf der Post-Neural Architecture Search vorgestellt.

In einer Reihe von Benchmark-Tests hat dieses Modell nicht nur eine vergleichbare oder sogar bessere Genauigkeit wie Qwen3, Qwen2.5, Gemma 3 und Llama 3.2 gezeigt, sondern auch eine bis zu 53,6-fache Beschleunigung der Generierungsdurchsatzleistung und eine 6,1-fache Beschleunigung in der Vorauffüllphase erreicht.

Es ist erwähnenswert, dass Jet-Nemotron-2B auf den MMLU-, MMLU-Pro- und BBH-Benchmarks einen 47-fach höheren Durchsatz als Qwen3-1.7B-Base erreicht hat und die Cache-Größe auf 1/47 reduziert wurde.

Zusätzlich hat es eine höhere Genauigkeit als DeepSeek-V3-Small und Moonlight (insgesamt 15 Milliarden Parameter, 2,2 Milliarden Aktivierungsparameter) erzielt.

Der Code und das vortrainierte Modell werden Open-Source veröffentlicht. Lassen Sie uns zunächst sehen, wie Jet-Nemotron aufgebaut ist.

Jet-Nemotron: Aufbau auf der Post-Neural Architecture Search

Zunächst wurde Jet-Nemotron auf der Grundlage der Post-Neural Architecture Search (PostNAS) entwickelt.

Die Post-Neural Architecture Search (PostNAS) ist eine Architektursuchmethode, die auf bestehenden großen Modellen aufbaut und diese verbessert.

Es geht von einem vortrainierten Voll-Attention-Modell aus und übernimmt direkt die Gewichte des Multi-Layer-Perceptrons, wobei diese Gewichte während des gesamten Prozesses eingefroren (nicht mehr aktualisiert) bleiben.

Jet-Nemotron wurde durch die folgenden 4 Schritte aus der PostNAS optimiert:

Platzierung und Entfernung von Voll-Attention-Schichten

Das Beibehalten einiger weniger Voll-Attention-Schichten im Modell ist für die Aufrechterhaltung einer hohen Genauigkeit bei anspruchsvollen Aufgaben wie der Suche von entscheidender Bedeutung.

Jedoch war die optimale Platzierung dieser Schichten bisher unklar.

Deshalb hat das Forschungsunternehmen eine neue Methode eingeführt, bei der ein "Once-for-All-Supernetzwerk" trainiert wird, um automatisch zu lernen, an welchen Positionen Voll-Attention-Schichten verwendet werden sollten.

Die Experimentergebnisse zeigen, dass diese lernende Platzierung im Vergleich zur üblichen gleichmäßigen Platzierungsstrategie eine signifikante Verbesserung der Genauigkeit auf dem MMLU-Benchmark erzielt.

Auswahl des linearen Attention-Moduls

Nach der Bestimmung der Platzierung der Voll-Attention-Schichten hat das Forschungsunternehmen eine Suche nach dem optimalen linearen Attention-Modul durchgeführt.

In den Experimenten wurden 6 der neuesten linearen Attention-Module evaluiert (RWKV7 wurde aufgrund des niedrigen Trainingsdurchsatzes ausgeschlossen). Die Ergebnisse sind wie folgt.

Aus der obigen Tabelle kann man entnehmen, dass Gated DeltaNet die beste Gesamtgenauigkeit erreicht hat. Deshalb wird Gated DeltaNet in den folgenden Experimenten verwendet.

Entwurf eines neuen Attention-Moduls

Das Hinzufügen von Faltungsschichten ist eine gängige Strategie zur Verbesserung der Fähigkeiten des linearen Attention-Mechanismus. Bisherige Methoden basierten jedoch nur auf statischen Faltungskernen und hatten keine Fähigkeit, die Merkmalsextraktionsmuster der Faltungskerne dynamisch anzupassen.

Deshalb hat das Forschungsunternehmen ein neues lineares Attention-Modul namens JetBlock eingeführt.

Dieses Modul verwendet einen Faltungskern-Generator, der dynamisch kausale Faltungskerne basierend auf dem Eingabeinhalt erzeugt und diese dann auf die V (Value)-Token anwendet. Darüber hinaus entfernt es die redundanten statischen Faltungsschichten auf Q (Query) und K (Key), um den Rechenprozess zu vereinfachen.

Durchführung einer hardwarebewussten Architektursuche

Traditionell wird die Anzahl der Parameter als Proxy für die Effizienz von Sprachmodellen verwendet. Jedoch ist die Anzahl der Parameter nicht direkt mit der Hardwareeffizienz korreliert.

Basierend auf der Erkenntnis, dass die Größe des KV-Caches der wichtigste Faktor für den Durchsatz bei langen Kontexten und langen Generierungen ist.

Das Forschungsunternehmen hat die Größe des KV-Caches auf die ursprüngliche Design-Spezifikation festgelegt und eine kleine Raster-Suche für die Key-Dimension, die Value-Dimension und die Anzahl der Attention-Heads durchgeführt.

Diese hardwarebewusste Suche kann bei gleichbleibendem Generierungsdurchsatz mehr Parameter nutzen, um eine höhere Genauigkeit zu erreichen.

Die gute Nachricht ist, dass das Forschungsunternehmen vorhat, den Code und das Modell auf GitHub zu veröffentlichen und derzeit auf die rechtliche Compliance-Prüfung wartet.

Deutliche Effizienzsteigerung

Jet-Nemotron-2B und Jet-Nemotron-4B basieren auf den Modellen Qwen2.5-1.5B bzw. Qwen2.5-3B.

Um die Leistung des Modells umfassend zu bewerten, hat das Forschungsunternehmen Tests in den Bereichen Mathematik, Allgemeinwissen, Suche, Codierung und Langkontext durchgeführt.

Bei mathematischen Aufgaben hat Jet-Nemotron-2B eine durchschnittliche Genauigkeit von 49,6 erreicht, was 6,3 höher als bei Qwen3-1.7B-Base ist und gleichzeitig 47-mal schneller ist.

Im Vergleich dazu lagen die früheren linearen Attention- und Mischmodelle bei mathematischen Aufgaben weit hinter Qwen3-1.7B-Base zurück.

Bei Allgemeinwissensinferenzaufgaben hat Jet-Nemotron-2B eine durchschnittliche Genauigkeit von 62,0 erreicht und alle Baseline-Modelle übertroffen.

Bei Suchaufgaben hat Jet-Nemotron-2B eine bessere Leistung als alle Baseline-Modelle außer Qwen3-1.7B-Base gezeigt.

Bei der Erweiterung auf 4B hat Jet-Nemotron-4B eine optimale durchschnittliche Genauigkeit von 76,2 erreicht und dennoch eine 21-fache Geschwindigkeitssteigerung im Vergleich zu Qwen3 beibehalten.

Bei Codierungsaufgaben hat Jet-Nemotron-2B eine höhere durchschnittliche Genauigkeit als alle Baseline-Modelle erreicht.

Zusätzlich hat Jet-Nemotron-4B in allen Codierungsaufgaben eine höhere Genauigkeit erzielt.

Bei Langkontextaufgaben kann man sehen, dass Jet-Nemotron-2B trotz nur zweier Voll-Attention-Schichten eine Leistung vergleichbar mit führenden Modellen wie Qwen2.5-1.5B und Gemma3n-E2B mit mehr Voll-Attention-Schichten erreicht.

Gesamtbetrachtet haben Jet-Nemotron-2B und Jet-Nemotron-4B in diesen Bereichen eine vergleichbare oder sogar bessere Leistung wie Qwen3-1.7B-Base gezeigt.

Und aufgrund der deutlichen Reduzierung der Voll-Attention-Schichten und der kleineren KV-Cache-Größe hat Jet-Nemotron einen deutlichen Vorteil gegenüber Qwen3.

Über das Team

Es ist erwähnenswert, dass alle Mitglieder dieses Forschungsteams Chinesen sind.

Yuxian Gu hat sein Bachelor- und Promotionsstudium an der Fakultät für Informatik und Technologie der Tsinghua-Universität absolviert. Sein Betreuer war Professor Huang Minlie.

Zuvor hat er auch ein Praktikum am Microsoft Research Asia absolviert, wo sein Betreuer der Forscher Dong Li war.

Seine Forschungsinteressen konzentrieren sich hauptsächlich auf den gesamten Lebenszyklus von Sprachmodellen, einschließlich des Pre-Trainings, der Anpassung an Downstream-Aufgaben und effizienter Methoden in der Inferenzphase.

Neuere Forschungsschwerpunkte sind die Theorie und Algorithmen für das Datenaufbauen von Pre-Trained Large Language Models (z. B. PDS, Instruction Pre-Training, Learning Law) sowie die Kompression von Sprachmodellen durch Wissensdistillation (z. B. MiniLLM, MiniPLM).