zero_grad 가 필요한 이유

학습 루프에서 자주 빠뜨리는 게 optimizer.zero_grad()다. 그런데 이게 없으면 학습이 망가진다. 왜 매 스텝마다 기울기를 초기화해야 하는지 짧게 정리해봤다.

loss.backward()는 기울기를 누적한다#

핵심은 loss.backward()의 동작이다. loss.backward()는 새로운 기울기를 재할당하는 게 아니라, 계속 += 방식으로 기존 기울기에 더해서 누적한다.

어디에 누적되나#

기울기가 쌓이는 곳은 파라미터 텐서 자신이 들고 있는 .grad 속성이다. 별도의 전역 저장소가 아니라 w.grad, b.grad 처럼 각 텐서에 붙어 있는 grad 버퍼다. loss.backward()는 연산 그래프를 거꾸로 타면서 requires_grad=True 인 leaf 텐서(보통 모델 파라미터)의 .grad에 계산한 기울기를 더해 넣는다. 중간 계산 결과 텐서는 미분이 계산되긴 해도 .grad에는 남지 않는다.

그래서 optimizer.step()은 이 .grad를 읽어서 파라미터를 갱신하고, optimizer.zero_grad()는 이 .grad를 0으로 비운다. 셋 다 같은 .grad를 공유한다.

같은 파라미터에 backward를 두 번 하면 .grad가 쌓이는 걸 그대로 볼 수 있다.

w = torch.tensor([2.0], requires_grad=True)

loss1 = (w ** 2).sum()   # dloss/dw = 2w = 4
loss1.backward()
print(w.grad)            # tensor([4.])

loss2 = (w ** 2).sum()   # 또 4
loss2.backward()
print(w.grad)            # tensor([8.])  ← 덮어쓰는 게 아니라 4+4 누적

왜 덮어쓰지 않고 누적하나#

이건 버그가 아니라 PyTorch의 의도된 설계다. 여러 경로에서 흘러온 기울기를 자연스럽게 합치기 위해서다.

• 한 파라미터가 여러 번 쓰이는 경우(가중치 공유 등)에는 각 경로의 기울기를 더해야 올바른 전체 기울기가 나온다. += 누적이라 별도 처리 없이 자동으로 합쳐진다.
• gradient accumulation 기법처럼 작은 batch로 backward만 여러 번 하고 step은 가끔 한 번씩 하는 방식도 누적이 기본 동작이라 그냥 된다.

즉 여러 기울기를 합치는 게 정상 시나리오라서 누적이 기본값이고, 매 스텝 독립적으로 학습할 때는 우리가 zero_grad()로 직접 끊어줘야 한다. 안 그러면 이전 스텝의 기울기가 .grad에 그대로 남은 채 이번 기울기가 더해져서 엉뚱한 방향으로 갱신된다.

학습 루프에서의 순서#

for epoch in range(epochs):
    optimizer.zero_grad()      # 누적된 기울기 초기화
    pred = model(x)            # 순전파
    loss = criterion(pred, t)
    loss.backward()            # 기울기 계산 (grad에 누적됨)
    optimizer.step()           # w, b 갱신

순서가 중요하다. backward로 기울기를 구하기 전에 zero_grad로 이전 기울기를 비워주고, 마지막에 step으로 갱신한다. zero_grad를 빼먹으면 backward의 += 누적 때문에 기울기가 스텝마다 부풀어 학습이 깨진다.