În ultimele două articole din această serie (Partea 1 și Partea 2), am explorat o gamă largă de trucuri arhitecturale și de antrenare pentru modelele de difuzie. Am încercat să evaluăm fiecare idee în izolare, măsurând debitul de date, viteza de convergență și calitatea finală a imaginii, încercând să înțelegem ce influențează cu adevărat performanța. În loc să optimizăm o singură dimensiune la un moment dat, în această a treia parte vom suprapune cele mai promițătoare ingrediente pentru a vedea cât de departe putem împinge performanța sub un buget strict de calcul. Acest lucru este foarte diferit de primele zile ale difuziei, unde antrenarea modelelor competitive putea costa milioane de dolari. Scopul aici este de a demonstra cât de mult a evoluat domeniul și cât de departe poate ajunge o inginerie atentă în doar o singură zi de antrenare.
Acest „speedrun” nu este doar un experiment amuzant. El va servi, cel mai probabil, ca fundație pentru rețeta noastră de antrenare la scară largă. Alături de rezultate, facem public codul sursă (link Github), o mișcare care demonstrează angajamentul nostru față de transparență și progresul comunității AI.
X-prediction și Antrenarea în Spațiul Pixelilor
O componentă fundamentală a acestei rețete rapide este utilizarea formulării x-prediction, inspirată din lucrarea „Back to Basics: Let Denoising Generative Models Denoise” [Li și He, 2025]. Așa cum am văzut în Partea a 2-a, acest lucru permite antrenarea direct în spațiul pixelilor și elimină complet necesitatea unui VAE (Variational Autoencoder). Aceasta este o schimbare de paradigmă semnificativă. În loc să depindem de un spațiu latent comprimat, care poate introduce artefacte și pierderi de informație, lucrăm direct cu datele brute.
Am utilizat o dimensiune a patch-ului de 32 și un bottleneck de 256 de dimensiuni în stratul inițial de proiecție a token-ilor. Acest design menține lungimea secvenței sub control, făcând antrenarea în spațiul pixelilor gestionabilă din punct de vedere computațional, chiar și la rezoluții mai mari. În loc să urmăm programul obișnuit 256px → 512px → 1024px, am început direct la 512px și apoi am făcut fine-tuning la 1024px. Cu un număr controlat de token-i și hardware modern, antrenarea în spațiul pixelilor nu mai este prohibitivă. Este pur și simplu o formulare mai curată și mai directă, care simplifică arhitectura generală a sistemului.
Pierderi Perceptive: Revenirea la Clasică
Un efect secundar foarte benefic al predicției x0 direct în spațiul pixelilor este că putem reutiliza o întreagă cutie de instrumente din viziunea computerizată clasică. Atunci când modelul produce latenți, supravegherea perceptivă devine dificilă. Trebuie fie să decodați înapoi la pixeli, fie să definiți pierderi într-un spațiu latent învățat care s-ar putea să nu se alinieze cu percepția umană. Odată ce preziceți pixelii direct, totul devine din nou simplu.
Ne-am inspirat din lucrarea „PixelGen: Pixel Diffusion Beats Latent Diffusion with Perceptual Loss” [Ma et al.], unde autorii introduc obiective perceptive suplimentare peste pierderea de difuzie. Ei demonstrează că adăugarea semnalelor perceptive poate îmbunătăți vizibil viteza de convergență și calitatea vizuală finală. Ideea este simplă: pe lângă obiectivul standard de potrivire a fluxului, încurajăm imaginea curată prezisă să se potrivească cu imaginea țintă într-un spațiu de caracteristici perceptive. LPIPS captează similaritatea perceptivă la nivel scăzut, în timp ce caracteristicile DINO oferă un semantic mai puternic.
În experimentele noastre, am constatat empiric că funcționează mai bine cu o greutate de 0,1 pentru pierderea LPIPS și 0,01 pentru pierderea perceptivă DINO. Aceste pierderi sunt ușoare comparativ cu trecerea principală a transformatorului, iar în configurația noastră adaugă doar un mic cost, oferind în același timp o creștere consistentă a calității.
Rutarea Token-ilor cu TREAD
Pentru a face fiecare pas mai ieftin din punct de vedere computațional, folosim rutarea token-ilor cu TREAD [Krause et al., 2025]. Această metodă selectează aleatoriu o fracțiune de token-i și îi lasă să ocolească un bloc contiguu de blocuri de transformator, reinjectându-i ulterior, astfel încât nimic nu este pierdut. Am ales TREAD în detrimentul SPRINT (Park et al., 2025) în principal pentru simplitate. Complexitatea suplimentară a SPRINT nu a justificat economiile de calcul destul de mici în configurația noastră.
Urmând rețeta TREAD, rutm 50% din token-i de la al 2-lea bloc până la penultimul bloc al transformatorului. Modelele rutate pot arăta mai rău sub CFG (Classifier-Free Guidance) vanilie, mai ales atunci când sunt sub-antrenate, așa că am implementat o schemă simplă de auto-ghidare inspirată din „Guiding Token-Sparse Diffusion Models”, care ghidează folosind o predicție condiționată densă versus rutată, în loc să se bazeze pe o ramură necondiționată.
Alinierea Reprezentării cu REPA și DINOv3
Am folosit REPA [Yu et al., 2024] pentru alinierea reprezentării. Pentru „profesor”, am ales DINOv3 [Siméoni et al. 2025], deoarece a oferit cele mai bune îmbunătățiri de calitate în experimentele noastre anterioare. Concret, aplicăm pierderea de aliniere o singură dată, la al 8-lea bloc transformator, cu o greutate a pierderii de 0,5. Deoarece combinăm REPA cu rutarea TREAD, calculăm pierderea de aliniere doar pe token-ii non-rutați, adică token-ii care trec efectiv prin blocurile unde aplicăm pierderea. Acest lucru menține semnalul REPA consistent și evită compararea caracteristicilor pentru token-ii care au sărit peste calea de calcul.
Optimizerul: Muon
Un alt element crucial a fost utilizarea optimizerului Muon, folosind implementarea FSDP din muon_fsdp_2, deoarece a arătat o îmbunătățire clară față de Adam în rulările noastre anterioare. Muon este aplicat doar parametrilor 2D (practic matrice). Tot restul (bias-uri, norme, embedding-uri etc.) este optimizat cu Adam, motiv pentru care configurația are două grupuri de parametri. Această abordare hibridă permite o optimizare mai fină și mai stabilă, adaptată tipurilor diferite de parametri din rețea.
Rezultate și Gânduri Finale
Pentru o rulare de antrenare de o singură zi, rezultatele sunt deja destul de solide. Modelul nu este încă perfect (încă se pot observa unele probleme de textură, anatomie ocazional ciudată și poate fi puțin instabil pe prompt-uri foarte dificile), dar este clar utilizabil. Urmărirea prompt-ului este puternică, estetica generală este consistentă, iar etapa de 1024 px face în principal ceea ce ne dorim: ascute detaliile fără a distruge compoziția.
Concluzia cheie este că suntem foarte aproape. Problemele rămase par a fi mai degrabă artefacte de sub-antrenare și diversitate limitată a datelor decât semne ale unui defect structural în rețetă. Modurile de eșec sunt consecvente cu ceea ce te-ai aștepta de la un model care pur și simplu nu a văzut încă suficiente date variate. Cu mai mult calcul și o acoperire mai largă, această configurație exactă ar trebui să continue să se îmbunătățească într-un mod destul de previzibil.
Privind în ansamblu, acest speedrun evidențiază cât de departe a ajuns antrenarea prin difuzie. Combinând antrenarea în spațiul pixelilor, rutarea eficientă, alinierea reprezentării și ghidarea perceptivă ușoară, puteți obține acum un model semnificativ în aproximativ o zi, cu un buget care ar fi sunat nerealist nu cu mult timp în urmă. Acest lucru democratizează accesul la tehnologii de vârf și deschide noi orizonturi pentru cercetare și inovare.
PRX Part 3 — Antrenarea unui model Text-to-Image în doar 24 de ore! O revoluție în ingineria modelelor de difuzie
