Loss Scaling

W krótkim formacie liczbowym bardzo mała wskazówka do poprawy modelu może zostać zaokrąglona do zera. Trening traci wtedy użyteczny sygnał. Chcesz zwiększyć wartości pośrednie, ale zachować właściwą wielkość końcowej aktualizacji.

Loss Scaling mnoży stratę — liczbową ocenę błędu — przez współczynnik przed obliczeniami wstecz. To zwiększa także gradienty, czyli wskazówki wpływu wag na błąd. Przed aktualizacją dzieli się je przez ten sam współczynnik.

W dokładnej arytmetyce obie operacje się znoszą. W ograniczonym formacie większa wartość może jednak przetrwać etap, w którym pierwotna zniknęłaby do zera. Odskalowanie musi odbyć się w odpowiedniej precyzji i kolejności.

Za duża skala grozi przepełnieniem, czyli liczbą poza zakresem formatu. Dlatego nie można dowolnie ją zwiększać. Skalowanie nie naprawia złych danych ani modelu; rozwiązuje konkretny problem reprezentacji liczb podczas treningu.

Mechanizm i szczegóły

$$\nabla_\theta(SL)=S\nabla_\theta L.$$

$L$ oznacza stratę, a $\theta$ parametry modelu. W arytmetyce dokładnej skalowanie i odskalowanie znoszą się. W ograniczonym formacie pośrednie, większe gradienty mogą przetrwać zapis, w którym ich pierwotne wartości zniknęłyby do zera. Ten mechanizm opisują Micikevicius et al., §3.2.

Okno zamiast dowolnie dużej skali

Własny przykład: $g=2^{-26}$ zaokrągla się w FP16 do zera. Po pomnożeniu przez 4 otrzymujemy $2^{-24}$, najmniejszą dodatnią wartość subnormalną FP16. Zapisujemy ją i dzielimy przez 4 w większej precyzji, odzyskując gradient. Pomnożenie gradientu, który już został zerem, niczego nie naprawi.

Jedna skala obejmuje jednak gradienty o różnych wielkościach. W demonstracji drugi wynosi 8. Skala $2^{13}$ daje dla niego 65536, czyli przepełnienie FP16. Zwiększając suwak, znajdź zakres, w którym oba gradienty przetrwają.

W Mixed Precision Training loss scaling chroni zakres gradientów; autocast dobiera formaty operacji. To osobne mechanizmy. Dokumentacja AMP, Gradient Scaling ostrzega, że FP16 nie pasuje do każdego modelu i skala może nawet spaść poniżej 1.

PyTorch GradScaler pomija krok optymalizatora, gdy wykryje nieskończone gradienty lub NaN. Przed gradient clipping trzeba odskalować gradienty, aby próg zachował znaczenie. Przy Gradient Accumulation skala pozostaje stała do zebrania całej efektywnej porcji; dopiero potem następuje odskalowanie i aktualizacja. Te granice pokazują przykłady AMP: unscaled gradients i accumulation.

Skalowanie nie zmienia celu Cross-entropy ani zamierzonego learning rate. Nie zastępuje też poprawnego uśrednienia straty po przykładach lub tokenach.