Skalowanie wnioskowania modeli dyfuzyjnych Nowy paradygmat

Wprowadzenie

Ostatnie postępy w dużych modelach językowych (LLM) wykazały skuteczność skalowania podczas wnioskowania. Modele takie jak o1, o3, DeepSeek R1, QwQ i Step Reasoner mini pokazały, że zwiększona moc obliczeniowa podczas wnioskowania może znacząco poprawić wydajność. Rodzi to pytanie: czy tę zasadę można zastosować również do modeli dyfuzyjnych? Zespół kierowany przez Xie Saining na Uniwersytecie Nowojorskim zbadał to właśnie zagadnienie. W wyniku systematycznych badań z wykorzystaniem ogólnych ram wyszukiwania stwierdzili, że skalowanie w czasie wnioskowania jest rzeczywiście skuteczne w przypadku modeli dyfuzyjnych. Zwiększenie mocy obliczeniowej podczas wnioskowania prowadzi do znaczącej poprawy jakości generowanych próbek. Co więcej, złożoność obrazów pozwala na różne kombinacje komponentów w ramach tych ram, dostosowane do różnych scenariuszy zastosowań.

Kluczowe Odkrycia

• Skalowanie czasu wnioskowania jest skuteczne dla modeli dyfuzyjnych: Alokowanie większej ilości zasobów obliczeniowych podczas wnioskowania skutkuje wyższą jakością próbek.
• Elastyczność kombinacji komponentów: Ramy pozwalają na różne konfiguracje komponentów, dostosowane do różnych zastosowań.
• Poza krokami odszumiania: Badania sugerują, że poszukiwanie lepszego szumu podczas próbkowania to kolejny wymiar skalowania NFE, poza zwykłym zwiększaniem liczby kroków odszumiania.
• Dwie osie projektowe: Ramy skupiają się na dwóch kluczowych osiach projektowych:
  • Weryfikatory: Dostarczanie informacji zwrotnej podczas procesu wyszukiwania.
  • Algorytmy: Znajdowanie lepszych kandydatów na szum.

Metodologia Badawcza

Zespół zbadał trzy różne scenariusze dla weryfikatorów, symulując różne przypadki użycia:

1. Scenariusze, w których dostępne są uprzywilejowane informacje na temat ostatecznej oceny.
2. Scenariusze, w których dostępne są informacje warunkowe, które kierują generowaniem.
3. Scenariusze, w których nie są dostępne żadne dodatkowe informacje.

W przypadku algorytmów badano:

1. Wyszukiwanie losowe: Wybieranie najlepszego z ustalonego zestawu kandydatów.
2. Wyszukiwanie zerowego rzędu: Iteracyjne ulepszanie kandydatów na szum za pomocą informacji zwrotnej od weryfikatora.
3. Wyszukiwanie ścieżki: Iteracyjne ulepszanie trajektorii próbkowania dyfuzyjnego za pomocą informacji zwrotnej od weryfikatora.

Badanie początkowo eksplorowało te projekty w stosunkowo prostym ustawieniu generowania warunkowego klasy ImageNet. Następnie zastosowano te projekty do generowania warunkowego tekstu na większą skalę i oceniono proponowane ramy.

Skalowanie Czasu Wnioskowania

Praca proponuje ramy skalowania czasu wnioskowania w modelach dyfuzyjnych, traktując to wyzwanie jako poszukiwanie optymalnego szumu próbkowania. Proces obejmuje dwa główne komponenty:

• Weryfikatory: Są to wstępnie wytrenowane modele, które oceniają jakość generowanych próbek. Pobierają one generowane próbki wraz z opcjonalnymi warunkami i zwracają wynik skalarny.
• Algorytmy: Te algorytmy wykorzystują wyniki weryfikatora do znalezienia lepszych próbek kandydujących. Funkcja pobiera weryfikator (V), wstępnie wytrenowany model dyfuzyjny (_θ) oraz zestaw wygenerowanych próbek i warunków, zwracając najlepszy szum początkowy.

Całkowity budżet wnioskowania jest mierzony całkowitą liczbą ocen funkcji (NFE), w tym zarówno kroków odszumiania, jak i kosztów wyszukiwania.

Weryfikatory Wyszukiwania

Badacze rozpoczęli od weryfikatora Oracle, który ma pełne informacje o ostatecznej ocenie wybranych próbek. W przypadku ImageNet obejmowało to metryki takie jak FID i IS. Następnie zbadali bardziej dostępne wstępnie wytrenowane modele jako nadzorowane weryfikatory, takie jak CLIP i DINO. Modele te były wykorzystywane do klasyfikowania próbek, wybierając próbkę z najwyższym logitem odpowiadającym etykiecie klasy.

Zaobserwowali jednak, że te klasyfikatory, działające punktowo, tylko częściowo pokrywają się z celami wyniku FID. Doprowadziło to do zmniejszenia wariancji próbek i załamania trybu wraz ze wzrostem obliczeń. Zjawisko to, nazwane 'hakowaniem weryfikatora', zostało przyspieszone przez nieograniczoną przestrzeń wyszukiwania algorytmu wyszukiwania losowego. Co ciekawe, badanie wykazało, że weryfikatory nie muszą koniecznie mieć informacji warunkowych, aby skutecznie kierować wyszukiwaniem. Zaobserwowano silną korelację między logitami z klasyfikatorów DINO/CLIP a podobieństwem cosinusowym przestrzeni cech między predykcją x na niskim poziomie szumu a ostateczną czystą próbką. Doprowadziło to do zastosowania samonadzorowanych weryfikatorów, które nie wymagają dodatkowych informacji warunkowych i nadal wykazywały skuteczne zachowanie skalowania.

Algorytmy Wyszukiwania

Aby złagodzić hakowanie weryfikatora, badacze zbadali bardziej wyrafinowane algorytmy wyszukiwania, które stopniowo optymalizują próbki kandydujące. Obejmowało to metodę wyszukiwania zerowego rzędu:

1. Rozpoczęcie od losowego szumu Gaussa jako punktu odniesienia.
2. Znalezienie N kandydatów w sąsiedztwie punktu odniesienia.
3. Przeprowadzenie kandydatów przez solver ODE w celu uzyskania próbek i wyników weryfikatora.
4. Aktualizacja punktu odniesienia o najlepszego kandydata i powtórzenie kroków 1-3.

Zbadano również algorytm wyszukiwania ścieżki, który bada możliwość wyszukiwania wzdłuż trajektorii próbkowania:

1. Próbkowanie N początkowych próbek szumu i uruchomienie solvera ODE do poziomu szumu σ.
2. Dodanie szumu do każdej próbki i symulacja procesu szumowania do przodu.
3. Uruchomienie solvera ODE na każdej zaszumionej próbce i zachowanie N najlepszych kandydatów na podstawie wyników weryfikatora, powtarzanie, aż solver ODE osiągnie σ=0.
4. Losowe przeszukiwanie pozostałych N próbek i zachowanie najlepszej.

Zarówno algorytmy wyszukiwania zerowego rzędu, jak i ścieżki zachowują silną lokalność w porównaniu z wyszukiwaniem losowym.

Skalowanie w Scenariuszach Tekst-Obraz

Zespół zbadał możliwości skalowania ram wyszukiwania w zadaniach tekst-obraz na większą skalę. Do oceny wykorzystano zbiory danych DrawBench i T2I-CompBench, a modelem bazowym był FLUX.1-dev. Rozszerzono również wybór nadzorowanych weryfikatorów, w tym Aesthetic Score Predictor, CLIPScore i ImageReward. Dodatkowo stworzono Verifier Ensemble, łącząc te trzy weryfikatory.

Analiza: Dopasowanie Weryfikator-Zadanie

Badanie porównało wyniki różnych kombinacji weryfikator-algorytm na różnych zbiorach danych. W DrawBench stwierdzono, że użycie wszystkich weryfikatorów ogólnie poprawiło jakość próbek. Zaobserwowano jednak, że użycie weryfikatorów Aesthetic i CLIP w izolacji może prowadzić do nadmiernego dopasowania ich uprzedzeń, co skutkuje negatywnym wpływem na siebie nawzajem. Wynika to z niedopasowania w ich fokusie oceny: Aesthetic Score skupia się na jakości wizualnej, często preferując wysoce stylizowane obrazy, podczas gdy CLIP priorytetyzuje dopasowanie wizualno-tekstowe, czasami poświęcając jakość wizualną. Zauważono, że niektóre weryfikatory są bardziej odpowiednie do konkretnych zadań, a skuteczność weryfikatora zależy od jego dopasowania do wymagań zadania.

Wydajność Algorytmu

Wszystkie trzy algorytmy wyszukiwania (losowe, zerowego rzędu i ścieżki) skutecznie poprawiły jakość próbkowania na DrawBench. Jednak wyszukiwanie losowe w niektórych aspektach wypadło lepiej ze względu na lokalny charakter pozostałych dwóch metod. Wyszukiwanie losowe szybciej zbiegało do uprzedzeń weryfikatora, podczas gdy pozostałe dwa algorytmy wymagają poprawy na mniej niż optymalnych kandydatach.

Kompatybilność z Dostrajaniem

Zespół zbadał kompatybilność swojej metody wyszukiwania z dostrojonymi modelami. Użyli modelu Stable Diffusion XL dostrojonego przez DPO i stwierdzili, że metoda wyszukiwania może być uogólniona na różne modele i poprawić wydajność już dopasowanych modeli.

Efekty Różnych Wymiarów Obliczeń Wnioskowania

Badanie zbadało, jak różne aspekty obliczeń wnioskowania wpływają na wyniki:

• Liczba Iteracji Wyszukiwania: Zwiększenie liczby iteracji przybliża szum do optimum.
• Obliczenia na Iterację Wyszukiwania: Dostosowanie liczby kroków odszumiania na iterację ujawnia różne regiony optymalne obliczeniowo.
• Obliczenia Generacji Końcowej: Zespół zastosował optymalne ustawienia dla końcowych kroków odszumiania, aby zapewnić najwyższą jakość próbki końcowej.

Efektywność Inwestycji w Obliczenia

Badacze zbadali skuteczność skalowania czasu wnioskowania na mniejszych modelach dyfuzyjnych. Stwierdzili, że w przypadku ImageNet skalowanie mniejszych modeli może być bardzo wydajne. W niektórych przypadkach wyszukiwanie na mniejszym modelu może przewyższyć większe modele bez wyszukiwania. Skuteczność zależy jednak od bazowej wydajności mniejszego modelu.

W ustawieniach opartych na tekście PixArt-Σ, wykorzystując tylko ułamek obliczeń, przewyższył FLUX-1.dev. Wyniki te pokazują, że znaczne zasoby obliczeniowe wydane podczas treningu mogą zostać zrównoważone mniejszymi ilościami obliczeń podczas generowania, co skutkuje wydajniejszymi próbkami wyższej jakości.