기울기 누적: 올바른 정규화를 유지하며 유효 배치 늘리기

정의

기울기 누적은 모델 파라미터를 갱신하기 전에 여러 마이크로배치를 처리한다. 이 기울기들을 한 번의 옵티마이저 스텝에 사용하므로, 물리적 배치가 장치 메모리에 들어가지 않아도 유효 배치를 크게 만들 수 있다. PyTorch의 backward 호출은 기존 파라미터 기울기에 결과를 더하며, 기울기를 지울 때까지 누적된다. [1][2]

작동 원리

마이크로배치 크기가 같고 손실이 동일 가중치 샘플의 평균이라면 각 손실을 누적 구간의 마이크로배치 수로 나눈다. 그러면 기울기 합은 전체 평균 손실의 기울기와 같다. 크기가 다르면 실제로 기여하는 샘플 수에 따라 가중치를 적용해야 한다. 토큰 손실은 마스킹되지 않은 유효 토큰을 세며, 클래스 가중치나 다른 축약 방식에서는 분모가 달라질 수 있다. 마이크로배치 평균을 무조건 평균 내면 목적함수가 바뀐다. 이 정규화 규칙은 원하는 가중 평균을 미분하면 도출된다.

실무 절차

구간 시작에 기울기를 지우고, 정규화한 각 마이크로배치의 순전파와 역전파를 실행한 뒤 한 번의 옵티마이저 스텝을 수행한다. 다음 구간을 위해 기울기를 지운다. 별도의 순전파에서는 보통 이전 계산 그래프를 유지할 필요가 없다. [2] 마지막 불완전 구간을 명시적으로 처리한다. 혼합 정밀도에서는 누적 중 스케일을 일정하게 유지하고 마지막 역전파 뒤 한 번만 스케일을 복원한다. 이후 합쳐진 기울기를 클리핑하고 업데이트하며, 스케일러도 이 경계에서 갱신한다. [1] 옵티마이저 업데이트 단위의 스케줄러는 실제 업데이트가 있을 때만 진행한다.

설명용 예시

한 장치가 샘플 여덟 개씩 든 마이크로배치 네 개를 처리한다고 하자. 유효 배치는 32개이므로 각 평균의 가중치는 1/4이다. 마지막 구간에 8, 8, 8, 2개의 샘플이 있으면 총 26개이며, 가중치는 8/26, 8/26, 8/26, 2/26이다. 이는 산술 예시이며 측정한 학습 결과가 아니다.

한계와 검증

배치 정규화는 물리적 마이크로배치별 통계를 사용하므로 누적이 하나의 큰 배치를 자동으로 재현하지는 않는다. [3] 무작위 연산과 부동소수점 덧셈 순서도 결과를 바꿀 수 있다. 마이크로배치마다 따로 클리핑하면 대개 합산 업데이트가 달라진다. 이 방식은 동시에 저장하는 활성화를 줄이지만 모델 가중치나 옵티마이저 상태를 없애지는 않으며, 추가 실행으로 시간이 늘어날 수 있다.

장시간 학습 전에는 배치 의존 층이 없는 작고 결정적인 모델에서 한 번의 누적 업데이트와 합친 배치 업데이트를 비교한다. 손실 정규화, 적절한 허용 오차 내 기울기 차이, 업데이트 횟수, 불완전 구간의 동작을 확인한다. 이는 권장 검증 항목이며 이 글을 위해 수행한 테스트가 아니다.

출처

[1] PyTorch: Automatic Mixed Precision examples [2] PyTorch: torch.Tensor.backward [3] PyTorch: BatchNorm1d