LeCuns Weltmodell kann auf einem einzigen GPU laufen.

Die neuesten Fortschritte des Weltmodells von LeCun: Es wurde ein minimalistisches Trainingsschema open source gemacht, das auf einer einzigen GPU läuft.

Dieses Schema heißt LeWorldModel. Es basiert auf der JEPA-Architektur und kann direkt aus Pixel-Eingaben die Zukunft vorhersagen. Die Geschwindigkeit ist erstaunlich schnell, eine vollständige Planung benötigt nur 1 Sekunde.

Es kann nur anhand von Pixelbildern, ohne komplizierte Techniken, auf einer einzigen GPU stabil trainiert werden und lernt, vorherzusagen, "Was passiert mit der Welt, wenn ich diese Aktion ausführe". Es eignet sich hervorragend zur Planung und Steuerung von Robotern und Agenten, ist schnell, stabil und einfach zu bedienen.

Die praktischen Ergebnisse sind auch sehr gut:

Hohe Geschwindigkeit: Die Planungsgeschwindigkeit ist 48-mal schneller als bei Lösungen mit großen Modellen und wird in weniger als 1 Sekunde abgeschlossen.

Wenige Parameter: Es hat nur 15 Millionen Parameter. Alle Trainings- und Planungsexperimente wurden auf einer einzigen NVIDIA L40S-Grafikkarte durchgeführt und können in wenigen Stunden abgeschlossen werden.

Starke Steuerung: In 2D/3D-Aufgaben wie Kisten schieben, Roboterarmen und Navigation übertrifft es die früheren End-to-End-Methoden und kann mit Lösungen mit großen Modellen mithalten.

Physikverständnis: Im Latent-Space sind physikalische Informationen wie Position und Winkel versteckt. Es kann auch "unphysikalische" Phänomene erkennen (z. B. wenn ein Objekt plötzlich teleportiert wird, wird es "überrascht" sein).

Technische Architektur: Die JEPA auf das Wesentliche reduzieren

Das Team erklärt, dass frühere JEPA-Methoden heuristische Techniken oder Tricks (z. B. EMA, Stopping-Gradient-Methode, vortrainierte Repräsentationen, Maskierung oder komplizierte Verlustfunktionen) anwendeten, um ein Kollabieren des Modells zu vermeiden.

Allerdings haben diese Techniken die JEPA-Trainingsinstabilität erhöht und die Durchführung erschwert.

Die Idee von LeWM ist es, die JEPA auf das Wesentliche zu reduzieren: Ein Encoder wandelt Bilder in Merkmale um, ein Prädiktor prognostiziert das nächste Merkmal basierend auf der Aktion, und eine Gaußsche Regularisierung verhindert das Kollabieren des Modells. Der gesamte Prozess ist End-to-End und minimalistisch und stabil.

Die Architektur verwendet nur zwei Kernkomponenten - Encoder + Prädiktor:

Encoder: Komprimiert das Bild in eine kleine Sequenz von Zahlen (Latent-Merkmale).

Prädiktor: Prognostiziert das nächste Merkmal basierend auf dem aktuellen Merkmal und der durchzuführenden Aktion.

Das wichtigste Innovation liegt darin, dass es nur zwei Verluste verwendet:

Vorhersageverlust: Der Prädiktor soll das reale Merkmal des nächsten Frames so genau wie möglich vorhersagen, indem er den einfachen MSE (Mean Squared Error) verwendet. Seine Aufgabe ist es, dem Modell die dynamischen Gesetze der Welt beizubringen.

SIGReg-Regularisierungsverlust: Erzwingt, dass alle Merkmalsvektoren einer Standard-Gauß-Verteilung folgen. Seine Aufgabe ist es, zu verhindern, dass das Modell "kollabiert" (alle Bilder geben die gleichen Merkmale aus).

Der endgültige Gesamtverlust = Vorhersageverlust + λ × SIGReg-Regularisierungsverlust.

Das Regularisierungsgewicht λ ist der einzige Hyperparameter, der optimiert werden muss. Dies vereinfacht den Trainingsablauf erheblich und macht LeWM stabil und einfach zu verwenden.

Experimentelle Ergebnisse: Übertrifft frühere JEPA-Methoden

Zunächst das Fazit: LeWM übertrifft die früheren End-to-End-JEPA-Methoden (PLDM) und ist mit dem auf großen Modellen vortrainierten DINO-WM gleichwertig oder sogar besser. Gleichzeitig ist es leichter zu trainieren, schneller und hat weniger Parameter.

Das Team hat Tests an 4 klassischen Robotik-/Steuerungsaufgaben durchgeführt und mit den beiden fortschrittlichsten JEPA-basierten Methoden DINO-WM und PLDM verglichen.

Die 4 Aufgaben sind Push-T (Kiste schieben), Reacher (Roboterarm greift nach einem Ziel), OGBench-Cube (3D-Roboterarm greift einen Würfel) und Two-Room (2D-Navigation).

Die Ergebnisse zeigen:

Push-T (Kiste schieben): LeWM ist am stärksten, mit einer Erfolgsrate von 96%, 18% höher als PLDM und sogar besser als DINO-WM mit sensorischem Input.

Reacher (Roboterarm greift nach einem Ziel): LeWM > PLDM und nahe an DINO-WM.

OGBench-Cube (3D-Roboterarm greift einen Würfel): LeWM ist etwas schwächer als DINO-WM, aber dennoch stark.

Two-Room (2D-Navigation): LeWM ist etwas schwächer, lernt aber die physikalischen Informationen dennoch gut.

In 2D- und 3D-Aufgaben verringert LeWM die Lücke zu Weltmodellen auf Basis von Basismodellen (z. B. DINO-WM) und übertrifft die End-to-End-Baseline PLDM.

Bemerkenswert ist, dass die Planungsgeschwindigkeit von LeWM 48-mal schneller ist als die von DINO-WM: Weniger als 1 Sekunde gegenüber etwa 47 Sekunden.

Der Grund ist, dass LeWM die Beobachtungsdaten um etwa 200-mal verkleinern kann, sodass die AI bei der Vorhersage der Zukunft schneller und mit weniger Aufwand rechnen kann und die auf Merkmalen basierte Weltplanung fast in Echtzeit durchgeführt werden kann.

Darüber hinaus versteht LeWM wirklich die Physik.

Das Modell wandelt das Bild in eine Sequenz von Zahlen (Latent) um. Das Team hat hinter dem trainierten LeWM einen einfachen Detektor angeschlossen, der nur anhand der Latent-Zahlen die Position des Roboters/Würfels, den Winkel des Würfels und die Koordinaten der Roboterarmspitze vorhersagen soll.

Die Positionsvorhersage ist fast 100% genau, und die Winkelvorhersage ist auch sehr genau. Es ist viel besser als das frühere PLDM und ähnlich gut wie das große Modell DINO.

Um die Lernleistung von LeWM anschaulich zu zeigen, hat das Team zusätzlich einen kleinen Dekoder für die Visualisierung trainiert, der drei Arten von Bildern zeigt: Echtzeitvideos, vom Modell "gesehene" rekonstruierte Videos und zukünftige Vorhersagevideos des Modells.

Man kann sehen, dass LeWM nicht nur die aktuelle Szene genau verstehen kann, sondern auch die zukünftige Bewegung der Objekte korrekt vorhersagen kann und wirklich die Kernstruktur und die Veränderungsgesetze der Umgebung erfasst.

Darüber hinaus kann es "physikalisch unregelmäßige" Phänomene erkennen.

Das Team hat ein Experiment durchgeführt und absichtlich zwei Arten von "Störungsszenarien" erzeugt, um zu sehen, ob das Modell etwas merkt:

Visuelle Störung: Ein Objekt ändert plötzlich seine Farbe.

Physikalische Störung: Ein Objekt teleportiert sich direkt an eine zufällige Position, was gegen die physikalischen Gesetze verstößt.

Das Modell reagiert auf die "Farbänderung" eher gleichgültig, während es bei der "physikalischen Verletzung" stark überrascht ist.

Team-Hintergrund

Der erste Autor Lucas Maes ist ein dritter Jahrgangsdoktorand am kanadischen AI-Institut Mila. Sein Betreuer ist Damien Scieur.

Derzeit ist er Gastforschungsassistent an der Brown University und arbeitet mit Randall Balestriero an Weltmodellen zusammen.

Seine Arbeit konzentriert sich auf die Verbesserung von JEPA durch verschiedene Methoden, einschließlich gradientenbasierter Planung, hierarchischer zeitlicher Abstraktion, Zieldefinition und physikalischem Verständnis.

Quentin Le Lidec ist ein Postdoktorand am Courant Institute of Mathematical Sciences der New York University und arbeitet mit Yann LeCun an Robotik-Weltmodellen zusammen.

Seine aktuelle Forschung konzentriert sich auf die Lösung von Problemen in der realen Welt mit Künstlicher Intelligenz. Er hat zu Open-Source-Projekten wie Pinocchio, Simple und stable-worldmodel beigetragen.

Damien Scieur ist derzeit Mitarbeiter bei Samsung und war früher Postdoktorand an der Princeton University. Sein Hauptgebiet sind Optimierungsalgorithmen.