Scaling Muon for Diffusion Transformers
Paper • 2608.20818 • Published
HomeIMG-1.2B-PL-epoch22
Dwujęzyczny (polsko-angielski) generator obrazów 256×256 trenowany od zera na 8 domowych kartach graficznych: denoiser DiT 1,24B warunkowany zamrożoną Poziomką-Bazą (polsko-angielski model językowy MoE, 4B) przez trenowalny refiner 72,1M. Rozumie prompty po polsku i po angielsku; FD-DINOv2 na polskich promptach 26,35 (poprzednik: 29,29).
To HomeIMG-641M-PL-epoch19 powiększony z 6 do 12 bloków i trenowany dalej przez 3 epoki. Nowe bloki dopisano na końcu z wyzerowanymi projekcjami wyjściowymi, jak w SANA 1.5, więc na starcie model liczył dokładnie to samo co poprzednik.
„epoch22”: denoiser trenował łącznie 22 epoki: 15 po angielsku ze Slayer149, 4 dwujęzyczne z Poziomką (6 bloków) i 3 dwujęzyczne po powiększeniu (12 bloków).
| po polsku | po angielsku |
|---|---|
 |
 |
Przykłady
Osiem promptów z poprzednich wydań, po polsku i po angielsku; każdy obrazek z seedem 0, więc para k ma ten sam szum. 50 kroków Eulera, CFG 3,0, wagi EMA. Kolejność wierszami:
| # | po polsku | po angielsku |
|---|---|---|
| 1 | rudy lis siedzący w zaśnieżonym lesie o wschodzie słońca, fotografia | a red fox sitting in a snowy forest at sunrise, photograph |
| 2 | futurystyczna panorama miasta nocą z neonami odbijającymi się w deszczu | a futuristic city skyline at night with neon lights reflected in the rain |
| 3 | miska świeżych truskawek na drewnianym stole, miękkie naturalne światło | a bowl of fresh strawberries on a wooden table, soft natural light |
| 4 | akwarela przedstawiająca latarnię morską na skalistym wybrzeżu | a watercolor painting of a lighthouse on a rocky coast |
| 5 | mały robot czytający książkę w przytulnej bibliotece | a small robot reading a book in a cozy library |
| 6 | astronauta jadący konno po Księżycu, grafika cyfrowa | an astronaut riding a horse on the moon, digital art |
| 7 | szyld z neonowym napisem OPEN | a street sign that says OPEN in neon letters |
| 8 | dwoje dzieci puszczających latawiec na zielonym wzgórzu pod błękitnym niebem | two children flying a kite on a green hill under a blue sky |
Wyniki
FD-DINOv2 (Stein i in. 2023, niżej znaczy lepiej): odległość Frécheta między
cechami CLS DINOv2 obrazów wygenerowanych i prawdziwych. 4096 odłożonych captionów (pierwszy caption pierwszych 4096
obrazów z 4 odłożonych shardów), 50 kroków Eulera, CFG 3,0, stałe seedy, cechy vit_base_patch14_reg4_dinov2.lvd142m;
referencja to prawdziwe odłożone latenty zdekodowane tym samym dekoderem FLUX.2. Kolumna PL: te same obrazy, ten sam
szum, prompty to polskie tłumaczenia tych samych captionów. Wartości nie są porównywalne z FD-DINOv2 z publikacji.
val/loss: velocity MSE na tych samych 4096 obrazach (angielskie captiony, stały szum i t).
| model | bloki | val/loss | FD-DINOv2 (EN) | FD-DINOv2 (PL) |
|---|---|---|---|---|
| HomeIMG-641M-epoch15 (WSM, Slayer149) | 6 | 0,6297 | 27,66 | – |
| HomeIMG-641M-PL-epoch19 (WSM) | 6 | 0,6277 | 28,90 | 29,29 |
| HomeIMG-1.2B-PL-epoch22 (to repo, EMA) | 12 | 0,6170 | – | 26,35 |
Najlepsza val/loss i najlepszy FD-DINOv2 ze wszystkich wydań HomeIMG: na polskich promptach 26,35 wobec 29,29 u poprzednika (−10%), czyli lepiej niż angielski HomeIMG-641M-epoch15 na promptach angielskich (27,66). Val/loss spadała równo do ostatniego kroku (0,6171 na kroku 8704, 0,6170 na kroku 8750), więc model nie jest nasycony.
Architektura
- Denoiser (1235,0M): DiT, 12 bloków o szerokości 2560, 20 głowic (128), SwiGLU 6912, jeden token na piksel latentu (siatka 16×16). Każdy blok: self-attention z RoPE 2D i QK-norm, cross-attention do tekstu (szerokość tekstu 1280), MLP SwiGLU; czas przez adaLN-single w stylu PixArt. Bloki 0–5 pochodzą z HomeIMG-641M-PL-epoch19, bloki 6–11 są nowe.
- Text encoder:
- Poziomka-Baza, zamrożona: BailingMoeV2, 16 warstw o szerokości 2048, 128 ekspertów; 8 z 32 aktywnych ekspertów na token (ok. 450M aktywnych parametrów; BLiMP 68,3% wobec 68,5% przy 32). Wagi ekspertów w int8 (skala na eksperta i kolumnę, mnożenia w bf16). Cechy tekstu: wyjście ostatniej warstwy.
- Refiner (72,1M): projekcja 2048 → 1280, 4 dwukierunkowe bloki transformera (szerokość 1280, 20 głowic q / 5 kv, RoPE 1D, SwiGLU 3456), RMSNorm, wyuczony pusty prompt do CFG. Do 255 tokenów promptu.
- Dekoder (zamrożony): FLUX.2 small decoder, latenty 128×16×16.
Trening
Start: wagi EMA etapu 2 HomeIMG-641M-PL-epoch19 (val/loss 0,6284), powiększone do 12 bloków. Nowe bloki 6–11 mają
wyzerowane attn.out, cross.out i mlp.out_proj, więc na starcie każdy z nich jest tożsamością. Nie kopiowano
istniejących bloków (w SANA 1.5 to się rozjeżdżało).
- Dane: QLNI/FLUX-Reason-6M-flux2-latents, 5,88 mln obrazów. Dla każdego obrazu losowy caption, w 50% przypadków po polsku (cpral/FLUX-Reason-6M-captions-pl, Index-Translate-2B; ~0,3% tłumaczeń oznaczonych jako wadliwe zastąpiono angielskim oryginałem).
- Długość: 3 epoki, 8750 kroków × 2016 = 17,6 mln próbek, trenowany cały model.
- Cel: rectified flow,
z_t = (1−t)·x + t·ε, celε − x, MSE, t ~ logit-normal(0, 1); 10% promptów zastąpionych pustym. Bez REPA. - Optymalizator: MuLoCo (7 replik uśrednianych co 64 kroki), Muon (CMuon, lr 3e-3, momentum 0,95 trzymane w bf16 ze stochastycznym zaokrąglaniem) na macierzach ukrytych, z Newtonem-Schulzem co 3. krok i normalizacją wierszy pomiędzy (periodic row-wise Muon); AdamW na reszcie. Rozgrzewka lr przez 2% kroków.
- Sprzęt: 8× RTX 3090 Ti, bf16. 7 kart trenuje (każda ma cały model, micro-batch 24, selektywny activation checkpointing), a ósma tylko liczy Poziomkę dla pozostałych, z wyprzedzeniem o krok: captiony są pakowane bez paddingu (varlen attention), a normy, router i attention Poziomki skompilowane. Newton-Schulz z własnym kernelem Triton, który liczy symetryczne iloczyny tylko w połowie (1,5× szybciej). Średnio ok. 285 obrazów/s, ok. 17 h.
- Wydane wagi: EMA (val/loss 0,61703). Najlepszy wariant WSM/QAM (równa średnia 3 ostatnich checkpointów) dał 0,61718, a surowe wagi z ostatniego kroku 0,61849.
Po treningu nowe bloki mają projekcje wyjściowe attn.out i mlp.out_proj o RMS ok. 30–40% starych bloków,
a cross.out 10–23% (mniej w głębszych blokach). Biorą już udział w obliczeniach, ale wciąż rosną.
Kod: https://github.com/adamo1139/trenowanie_generatora_obrazkow
Użycie
Repozytorium jest kompletne: zawiera Poziomkę (text_encoder/, bf16, eksperci w int8, dokładnie w tej postaci, na
której trenowano; 4,1 GB) i jej tokenizer. Potrzebne są diffusers z AutoencoderKLFlux2 i ok. 10 GB VRAM:
import torch
from diffusers import DiffusionPipeline
REPO = "cpral/HomeIMG-1.2B-PL-epoch22"
pipe = DiffusionPipeline.from_pretrained(REPO, custom_pipeline=REPO, trust_remote_code=True).to("cuda")
image = pipe("rudy kot śpiący na parapecie, fotografia", num_inference_steps=50, guidance_scale=3.0,
generator=torch.Generator("cuda").manual_seed(0)).images[0]
image.save("kot.png")
Prompty po polsku albo po angielsku, do 255 tokenów. Obraz zawsze ma 256×256.
Ograniczenia
- Model badawczy: tylko 256×256, ok. 6 mln syntetycznych obrazów (FLUX).
- Nowe bloki są niedotrenowane: po 3 epokach mają mniejsze wagi niż stare, model zyskałby na dłuższym treningu.
- Twarze bywają zniekształcone, zwłaszcza w zbliżeniach i małe postacie w kadrze.
- Rzadsze polskie słowa są czasem rozumiane źle; polskie captiony to tłumaczenia maszynowe modelu 2B.
- Słabe strony: renderowanie tekstu, dokładność kompozycji, uprzedzenia danych treningowych.
Licencja
Refiner, denoiser i kod: Apache-2.0. text_encoder/ zawiera też wagi
Poziomki-Bazy (zamrożone, eksperci przekonwertowani do int8)
na jej licencji, openmdw-1.1. vae/: FLUX.2 small decoder © Black Forest Labs, Apache-2.0, bez zmian. Dane mają
własne warunki, zobacz karty zbiorów danych.
- Downloads last month
- -
Model tree for cpral/HomeIMG-1.2B-PL-epoch22
Base model
cpral/HomeIMG-641M-epoch15 Finetuned
cpral/HomeIMG-641M-PL-epoch19