Wizualizacja funkcji straty w 3D: co naprawdę pokazuje wykres

Czym właściwie jest wizualizacja funkcji straty w 3D?

Funkcja straty przypisuje wartość każdemu ustawieniu parametrów modelu, a celem optymalizacji jest znalezienie ustawienia, które tę wartość minimalizuje [1]. Sieć neuronowa ma tysiące, miliony albo więcej parametrów, więc pełnej funkcji nie da się narysować. Wizualizacja 3D wybiera dwa kierunki w przestrzeni parametrów i pokazuje przekrój zapisany jako L(θ+aδ+bη), gdzie wartość straty jest trzecią osią albo kolorem [2].

Gradient wskazuje lokalny kierunek największego wzrostu funkcji, dlatego zejście gradientowe przesuwa parametry w stronę przeciwną do gradientu [1]. Na wykresie widać wtedy trajektorię: kolejne kroki optymalizatora wędrują po powierzchni rozpiętej na dwóch wybranych kierunkach. Kolor punktu opisuje wartość straty dla konkretnego ustawienia parametrów, a kształt powierzchni opisuje, jak ta wartość zmienia się w tych dwóch kierunkach. O pozostałych kierunkach wykres milczy [2].

Zbyt duży learning rate powoduje przeskakiwanie minimum albo rozbieżność, a nie gładki zjazd do dna doliny [1]. Na przekroju wygląda to jak zygzak między ścianami wąskiego zagłębienia. Ten sam objaw może jednak pochodzić z innej przyczyny, na przykład z batch normalization, która zmienia kształt powierzchni wokół wag [2].

Czy wykres 3D pokazuje całą funkcję straty sieci?

Nie. To projekcja, nie mapa. Li, Xu, Taylor, Studer i Goldstein opisali w pracy z NeurIPS 2018 metodę rzutowania krajobrazu straty na dwa kierunki [3]. Pokazali, że surowe wykresy zniekształcają ocenę ostrości minimum i zaproponowali normalizację filtrów, żeby porównania między modelami były wiarygodne [2].

Bez normalizacji model o większej skali wag wygląda na ostrzejszy tylko dlatego, że ma większe wagi. To prowadzi do błędnych wniosków o generalizacji, bo porównuje się modele w różnych skalach, a nie różne geometrie [2]. Autorzy opisali też, że liczba kierunków i sposób ich losowania zmieniają obraz bardziej niż sama zmiana optymalizatora [3].

Co przekrój mówi o SGD i Adamie?

SGD z momentum porusza się po powierzchni z bezwładnością. Przy zbyt dużym kroku przeskakuje płytkie zagłębienia albo oscyluje między ścianami doliny [1]. Adam skaluje krok osobno dla każdego parametru, korzystając z wykładniczych średnich gradientu i jego kwadratu oraz z korekcji obciążenia na początku treningu [4]. Na przekroju tor Adama jest krótszy i gęstszy w okolicy startu, bo algorytm szybko wchodzi w najbliższe zagłębienie.

Trzy wiersze tabeli pokazują, że różnica nie leży w jednej liczbie, ale w trzech decyzjach projektowych naraz: skali kroku, pamięci i zachowaniu na wąskim minimum. Sama zmiana optymalizatora bez zmiany learning rate daje na wykresie obraz, którego nie da się porównać z poprzednim [1].

Jak samodzielnie narysować krajobraz straty w PyTorch?

Potrzebny jest wytrenowany model, dwa kierunki w przestrzeni wag i siatka punktów wokół wag. Kod referencyjny z pracy Li i in. jest publicznie dostępny i pokazuje całą procedurę [7].

1. Wytrenuj model i zapisz state_dict.
2. Wylosuj dwa kierunki δ i η o rozkładzie normalnym i znormalizuj je tak jak w oryginalnym kodzie [7].
3. Policz stratę na siatce współczynników a i b wokół wag, na przykład 21 × 21 punktów.
4. Narysuj powierzchnię i pokoloruj ją wartością straty.

Pętla treningowa w PyTorch ma stałą kolejność: zero_grad(), predykcja i strata, loss.backward(), optimizer.step() [6]. Lista optymalizatorów i ich parametrów jest w dokumentacji torch.optim [5]. Wersja 2.14 opisuje między innymi foreach i capturable, które zmieniają sposób wykonania kroku, nie jego matematykę [5].

Do porównania trajektorii na jednym wykresie zapisz współrzędne wag w wybranych epokach i nanieś je na tę samą siatkę. Jeśli interesuje cię szerszy obraz wdrożenia modeli, zobacz tekst o praktycznych zastosowaniach AI.

Kiedy wizualizacja wprowadza w błąd?

Cztery rzeczy psują odczyt. Pierwsza: wybór kierunków, bo dwa losowe kierunki dają inny obraz niż dwa kierunki własne Hessianu [2]. Druga: skala osi, bo bez wspólnego zakresu dwóch modeli nie da się porównać [2]. Trzecia: kontekst architektury, bo batch normalization i dropout zmieniają kształt powierzchni [2]. Czwarta: punkt startowy, bo wykres wokół wytrenowanych wag nie mówi, jak wyglądała droga do nich.

Wysoka wymiarowość nie jest wyłącznie przeszkodą. W przestrzeniach o dużej liczbie wymiarów większość kierunków jest do siebie prawie prostopadła, a minima bywają połączone płaskimi ścieżkami. To zjawisko opisuje literatura o błogosławieństwie wymiarowości [8]. Wniosek z jednego wykresu 2D bywa więc zbyt pesymistyczny.

Jak wykorzystać wykres w diagnozie treningu?

Trzy sygnały są użyteczne. Oscylacje przy stabilnej średniej straty wskazują na zbyt duży learning rate [1]. Wąskie minimum przy dobrej stracie treningowej i słabej walidacyjnej sugeruje ostrą geometrię, którą zestawiasz z regularyzacją [2]. Płaski obszar o wysokiej stracie oznacza, że model utknął daleko od dobrego rozwiązania i problem leży w danych albo architekturze, nie w optymalizatorze [1].

Wykres nie jest dowodem. Zestaw go z krzywymi straty treningowej i walidacyjnej, normą gradientu i metryką na zbiorze testowym [2]. Dopiero zgodność kilku sygnałów daje podstawę do zmiany. Praktyczny przykład pracy z narzędziami AI w ograniczonym czasie opisuje tekst o Codexie w data science.

Najczęstsze pytania

Czy wizualizacja funkcji straty wymaga GPU?

Nie. Siatka 21 × 21 punktów na małym modelu liczy się na CPU w rozsądnym czasie. GPU przyspiesza obliczenia, gdy model jest duży albo siatka gęsta. Oryginalny kod z pracy Li i in. uruchamia się na pojedynczym GPU, ale zasada obliczeń jest ta sama [7].

Ile kierunków potrzeba do wiarygodnego przekroju?

Dwa kierunki to minimum i standard w pracach z 2018 roku [3]. Jeden przekrój bywa mylący, więc losowanie powtarza się kilka razy i porównuje obrazy. Trzy lub cztery niezależne przekroje dają pojęcie o zmienności kształtu, a nie o pełnej geometrii funkcji [2].

Czy Adam zawsze schodzi szybciej niż SGD?

Nie. Adam szybciej zmniejsza stratę na początku dzięki adaptacyjnym krokom i korekcji obciążenia [4]. Na dłuższych treningach SGD z momentum często osiąga lepszą generalizację, bo nie zatrzymuje się w pierwszym wąskim zagłębieniu [1]. Wybór zależy od zadania, a nie od samego wykresu.

Co oznacza ostre minimum?

Ostre minimum to zagłębienie, w którym mała zmiana parametrów mocno podnosi stratę. Taki punkt częściej źle się generalizuje, bo dane testowe przesuwają model poza wąską dolinę [2]. Normalizacja filtrów jest potrzebna, żeby porównywać ostrość między modelami o różnej skali wag [2].

Źródła

1. Loss Functions and Optimization (Stanford CS231n)
2. Visualizing the Loss Landscape of Neural Nets (Cornell, 2018)
3. Visualizing the Loss Landscape of Neural Nets (NeurIPS 2018)
4. Adam: A Method for Stochastic Optimization (ICLR 2015)
5. torch.optim — PyTorch 2.14 documentation
6. Optimizing Model Parameters — PyTorch Tutorials
7. tomgoldstein/loss-landscape — GitHub
8. Loss landscapes and the blessing of dimensionality (Medium)