Pyramid-JiT verzichtet auf VAE und schlägt das Vorgängermodell mit 11,3-mal weniger Trainingsdaten
Training Text-to-Image Models Without a VAE
Linum AI stellt Pyramid-JiT vor, eine neue Pixelraum-Architektur, die ohne VAE auskommt. Durch mehrstufige Vorhersagen entlang des DiT-Trunks erreicht das Modell die Qualität des vorherigen Latent-Diffusions-Modells Linum v2 mit 11,3-mal weniger Trainingsbeispielen und 4,3-mal weniger GPU-Stunden bei vierfacher Pixelzahl. Der Code und die Gewichte stehen unter Apache 2.0 zur Verfügung.
Der Verzicht auf die VAE erlaubte es uns, Tokens aggressiver zu komprimieren. Das senkte die Kosten der Attention und führte direkt zu schnellerem Training und schnellerer Inferenz.
- schopra909
Hi HN, hier der Autor!
Als Kontext: Wir sind ein 2-Personen-Labor, das generative Videomodelle trainiert. Ziel ist ein neues Set an steuerbaren Animationswerkzeugen (mehr dazu hier https://www.linum.ai/about, falls es euch interessiert).
Der größte Engpass für unser letztes Text-zu-Video-Modell in Bezug auf Trainings- und Inferenzkosten ist Attention. Videomodelle sind unglaublich token-dicht (z. B. 110K Tokens für einen mehrere Sekunden langen Clip). Wenn wir dieses Kontextfenster aggressiver verdichten können, können wir größere Modelle für viel weniger $$ trainieren und sie Prosumern zu vernünftigen Preisen anbieten (anders als die großen Modelle von heute wie Seedance, deren Betrieb ein Vermögen kostet).
Traditionell haben Bild- und Videomodelle zwei getrennte Komponenten: VAE (Variational Autoencoder) und Diffusion Transformer (DiT). Sie werden separat trainiert, und empirisch scheint der VAE Mühe zu haben, über eine 16x16-Token-Reduktion hinauszukommen.
Hier wechseln wir in den Pixelraum, werfen den VAE weg und erreichen eine 32x32-Token-Reduktion (4x kleinere Kontextfenster), während wir in einem Bruchteil der Trainingssamples ein insgesamt besseres Modell lernen.
Die zentrale These lautet "einfacher ist besser". Wenn wir das Kompressionsproblem in den leistungsfähigeren Diffusion Transformer (DiT) verlagern können, sollten wir in der Lage sein, einen für die Generierung optimierten "latenten Raum" zu lernen und eine bessere Kompression zu erreichen, ohne die Generierungsqualität zu beeinträchtigen.
Ich werde diesen Beitrag in den nächsten paar Stunden immer mal wieder checken, also stellt gerne Fragen unten. Und ich werde versuchen, sie so gut wie möglich zu beantworten […]
- vunderba
Schön. Was haltet ihr von dem kürzlich veröffentlichten Iris-3B, einem Pixelraum-Modell, das ebenfalls ohne traditionellen VAE auskommt?
- joefourier
Ich frage mich, wie viel davon auf die Verwendung eines suboptimalen VAE zurückzuführen ist? Dieselben Optimierungen könnten darauf angewendet werden, und meine Intuition sagt mir, dass euer gesamter Compute-Aufwand noch optimaler wäre, wenn ihr den VAE selbst mit einem besseren Ansatz neu trainieren würdet (insbesondere um Translations- und Rotationsinvarianz zu gewährleisten, + die Fähigkeit, die Latents neu zu skalieren/zu verwischen).
Ihr könntet auch dieselben Vorteile eines Entwurfs in einem niedrig aufgelösten latenten Raum haben, mit einer leichter zu erlernenden Datenverteilung, die robuster gegenüber Perturbationen ist, aber statt 512x512 könntet ihr bei gleichem Compute eine 4096x4096-Ausgabe erhalten (unter der Annahme eines 8x VAE).
Es gibt auch den Vorteil, eine hohe Anzahl von Diffusions-/Flow-Matching-Schritten für das Haupt-LDM verwenden zu können, während der VAE einstufig und viel kleiner sein kann, da er weder Sprache noch signifikantes Szenenverständnis bewältigen muss, sondern nur perzeptuelle Kompression. Das klingt besonders wichtig für ein Videomodell, bei dem ich extrem zögern würde, ein generatives Modell zu trainieren, ohne auf Interframe-Kompression zurückzugreifen.