Diese Präsentation stellt NNXv2 vor, eine grundlegende Überarbeitung des NNX-Frameworks für JAX, die in Kürze als Teil von Flax 0.13 veröffentlicht wird. Der Schwerpunkt liegt auf der Vereinfachung von APIs, der Durchsetzung reiner Baumstrukturen (PyTrees) anstelle von Graphen, der Entfernung von Graph-Updates als Standardverhalten und der engeren Angleichung an native JAX-Transformationen. Neue APIs wie unpack, jit_partial sowie nx.capture und so reduzieren den Boilerplate-Code und minimieren den Python-Overhead bei komplexen neuronalen Netzen.
Wichtigste Punkte & Zeitstempel
0:00 NNXv2 Ankündigung: Vorstellung von NNXv2 als neuer Version zur Vereinfachung von APIs und zur Umsetzung von Lektionen aus den Vorjahren.
1:15 Funktionsweise von NNX: NNX nutzt Python-Code zur Erstellung von PyTrees und bildet veränderliche Zustände (Mutabilität) wie Modellparameter oder Zähler ab.
4:18 Problem von Graph-Updates: Historisch unterstützte Graph-Updates erlaubten das dynamische Hinzufügen von Attributen im laufenden Betrieb, verursachten jedoch erhebliche interne Komplexität.
5:27 Übergang zum Tree-Modus: NNXv2 setzt standardmäßig strikte Baumstrukturen voraus (graph=false), wodurch automatische Transformationen drastisch vereinfacht und beschleunigt werden.
8:18 Variablen-Schließungen und Erfassung: Einführung nativer Unterstützung für Variable-Closures (Variablen-Schließungen) und neue Mechanismen zur Erfassung von Zwischenwerten.
9:33 Neue APIs unpack und jit_partial: Nutzung von unpack zur manuellen Trennung trainierbarer Parameter von Nicht-Parametern, sowie jit_partial zur Vorab-Flattung von verschachtelten PyTrees zur Minimierung des Python-Overheads.
11:55 Metrik-Erfassung mit so und nx.capture: Ersatz für Graph-Updates durch die so- und nx.capture-APIs zur transparenten und kontrollierten Weiterleitung von Zwischenwerten und Metriken aus verschachtelten Schichten.
13:40 Zukunftsausblick und Release: NNXv2 ist stark auf das zukünftige hijacks-Konzept ausgerichtet und erscheint als Breaking Change in Flax 0.13.
Glossar
NNX: Eine Erweiterungsbibliothek für das JAX-Ökosystem, die objektorientierte Programmierkonzepte für den Bau neuronaler Netze bereitstellt.
JAX: Ein von Google entwickeltes Python-Framework für hochleistungsfähiges numerisches Rechnen und automatisches Differenzieren, optimiert für GPUs und TPUs.
PyTree (Py-Baum): Eine baumartige Datenstruktur aus Python-Standardtypen (wie Listen, Tupeln und Dictionaries), die von JAX nativ verarbeitet und transformiert werden kann.
Mutabilität (Veränderbarkeit): Die Eigenschaft von Datenstrukturen oder Variablen, ihren Zustand nach der Erstellung während der Programmausführung zu verändern (z. B. aktualisierte Modellgewichte).
JIT (Just-In-Time-Kompilierung): Ein Verfahren, bei dem Programmcode zur Laufzeit in maschinennahen Code kompiliert wird, um die Ausführungsgeschwindigkeit drastisch zu steigern.
API (Programmierschnittstelle): Eine definierte Schnittstelle, die es verschiedenen Softwarekomponenten ermöglicht, miteinander zu kommunizieren.
Closure (Funktionsschließung): Ein Programmierkonzept, bei dem eine Funktion auf Variablen außerhalb ihres eigenen lokalen Gültigkeitsbereichs zugreift.
Flattung / Unflattung: Das Umwandeln verschachtelter Datenstrukturen (wie PyTrees) in eine lineare, flache Liste (und umgekehrt), was für bestimmte Optimierungsschritte in JAX erforderlich ist.
Flax: Ein populäres Open-Source-Framework für neuronale Netze, das auf JAX aufbaut.
Loss (Verlustfunktion): Eine mathematische Funktion, die misst, wie fehlerhaft die Vorhersagen eines Modells im Vergleich zu den tatsächlichen Zielwerten sind.