Adam

Model popełnia błędy i trzeba zmieniać jego wagi. Wskazówki z kolejnych grup danych wahają się, a różne wagi mogą otrzymywać wartości o różnej skali. Chcesz wykorzystać historię tych wskazówek, zamiast reagować wyłącznie na ostatnią.

Adam to optymalizator, czyli algorytm wyliczający aktualizacje wag. Wagi są liczbami sterującymi obliczeniami. Gradient mówi, jak ich mała zmiana wpływa na stratę — liczbową ocenę błędu.

Dla każdej wagi Adam pamięta średnią kroczącą gradientów oraz średnią ich kwadratów. Pierwsza podtrzymuje informację o kierunku, druga pomaga dostosować skalę kroku. Nowsze obserwacje mają większy udział; algorytm uwzględnia też korektę początkowego obciążenia tych średnich.

To ułatwia prowadzenie treningu, ale nie gwarantuje znalezienia najlepszego rozwiązania. Dwie zapamiętane liczby na wagę zajmują dodatkową pamięć, którą ZeRO może rozdzielać między urządzenia. Szczegóły obliczeń i ograniczenia opisano poniżej.

Mechanizm i szczegóły

Najpierw model przetwarza przykłady, a funkcja straty, np. Cross-entropy, mierzy błąd. Następnie obliczany jest gradient: informacja, jak mała zmiana każdej wagi wpływa na stratę. Adam korzysta z gradientów, żeby wyznaczyć aktualizację wag.

Co Adam pamięta?

Dla każdej wagi przechowuje dwie liczby:

• Średnią kroczącą gradientów — zachowuje informację o kierunku zmian z poprzednich kroków. Nowsze gradienty mają większy udział niż starsze.
• Średnią kroczącą kwadratów gradientów — opisuje ich skalę bez względu na znak. Służy do dostosowania wielkości aktualizacji osobno dla każdej wagi.

Są to estymaty pierwszego i drugiego momentu. Druga liczba nie jest zwykłą wariancją: algorytm uśrednia kwadraty gradientów bez odejmowania średniej. Adam koryguje też początkowe zaniżenie obu średnich wynikające z rozpoczęcia od zera. Mechanizm opisują Kingma i Ba, §2–3 i algorytm 1.

Mały przykład

Własne podstawienie do reguły Adama: po korekcie początkowego zaniżenia średnia gradientów wynosi 2, a średnia ich kwadratów 4. Adam dzieli 2 przez pierwiastek z 4, otrzymując 1. Przy współczynniku uczenia (learning rate) 0,001 odejmuje więc od wagi około 0,001. Waga 0,5 zmienia się na około 0,499. Pomijamy tu bardzo małą stałą dodawaną w mianowniku dla stabilności obliczeń.

Adam nadal wymaga dobrania współczynnika uczenia. Dostosowanie kroku do historii gradientów nie gwarantuje spadku błędu po każdej aktualizacji.

Dlaczego pojawia się przy ZeRO?

Dwie statystyki na każdą z miliarda wag, zapisane jako liczby FP32 po 4 bajty, zajmują 8 GB — jeszcze bez samych wag, gradientów i pozostałych danych treningu. To własny rachunek, z dziesiętnym GB.

ZeRO rozdziela między GPU m.in. te pomocnicze dane, czyli stan optymalizatora. Adam określa sposób aktualizacji wag, a ZeRO organizuje przechowywanie i wymianę potrzebnych danych. W recepturze Mixed Precision Training analizowanej w pracy ZeRO, §3.1 dochodzi jeszcze kopia wag FP32.

Zobacz także: AdamW — wariant, który oddziela kurczenie wag od kroku wynikającego z danych.