Pytorch로 딥러닝 훈련할 때 weight가 업데이트 되는 process

 오늘은 pytorch로 딥러닝 훈련을 할 때 weight가 어떻게 업데이트되는지 공부해보자. 

 

Recall : backpropagation 

 딥러닝에서 backpropagation + gradient descent로 학습한다는 것은, 현재 weight 값(예: a=1, b=3)에서의 loss에 대한 weight의 편미분 값을 구하고, 그 편미분 값(gradient)의 반대 방향으로 learning rate만큼 weight를 이동시키는 과정을 batch마다 반복하는 것이라고 recall할 수 있다. 

<현재 weight 값에서의 loss에 대한 편미분>

<weight update식>

 

Preliminaries 

A) preliminary : pytorch tensor 

pytorch tensor는 다음 두 가지를 함께 가진다.

1. 텐서의 값(data/value): 실제 파라미터 값 (예: weight 행렬의 원소들)
2. 텐서의 기울기(grad): 해당 텐서의 값에서 다른 텐서에 대한 해당 텐서의 편미분값을 저장하는 곳 
   
   • 내가 아무것도 안했는데 채워지는건 아니고, 이 텐서가 loss의 계산에 대해서 기여를 했다는 가정 하에, loss.backward()를 호출했을 때 해당 텐서의 값에서 loss에 대한 해당 텐서의 편미분 값이 채워진다. (뒤에 코드를 보면 더 잘 이해될 것이다) 

B) preliminary : optimizer

optimizer가 하는 일은 개념적으로 단순하다. optimizer에 등록된 파라미터(weight)들을 순회하면서, 각 파라미터에 저장된 gradient(grad)를 읽어 현재 파라미터 값(value)을 업데이트한다.

 가장 기본적인 형태에서는 gradient에 learning rate를 곱해, 그 값을 파라미터에서 빼는 방식으로 업데이트가 이루어진다.

Adam과 같은 optimizer들은 여기에 더해, 과거 gradient 정보를 이용한 momentum(1차 모멘트) 및 분산 추정(2차 모멘트) 등을 계산하여, 단순한 gradient descent보다 안정적이고 효율적인 업데이트를 수행한다.

핵심은 optimizer의 종류와 상관없이 동일하다. optimizer는 각 파라미터에 저장된 grad 값을 기반으로, 자신이 정의한 규칙에 따라 파라미터의 실제 값(value)을 갱신한다.

코드 분석

이제 모든 준비는 끝났다. 코드를 분석해보자.  dataloader는 잘 짜여져있다고 가정했을 때, 딥러닝 모델의 훈련 코드는 다음과 같다. 

for batch in dataloader:
    optimizer.zero_grad()   # 1. 이전 grad 초기화
    output = model(batch)   # 2. forward
    loss = criterion(output)
    loss.backward()         # 3. backward
    optimizer.step()        # 4. weight update

 

1) 데이터를 가져오고 forward로 loss를 만든다

• DataLoader가 배치 데이터를 준다.
• 모델이 그 배치를 받아 forward를 수행하고, 최종적으로 loss를 계산한다.

여기서 중요한 점은, PyTorch의 autograd가 loss가 만들어지기까지 거쳐온 연산 그래프(computation graph) 를 내부적으로 기억한다는 것이다. 즉 “loss를 만드는 데 관여한 텐서/연산”들이 연결된 그래프가 구성된다.

2) loss.backward() : loss에 대한 각 파라미터의 gradient 계산

loss.backward()를 호출하면,

• loss를 만들기까지 사용된 연산 그래프를 따라가며(역전파),
• 각 파라미터(leaf tensor, nn.Parameter)의 loss에 대한 편미분값을 계산하고,
• 그 결과를 각 파라미터 텐서의 param.grad에 저장한다.

3) optimizer.step() : 등록된 파라미터를 grad 방향으로 업데이트

optimizer.step()을 하면,

• optimizer에 등록된 파라미터 목록(예: model.parameters())에 대해,
• 각 파라미터의 .grad를 읽어서
• 학습률(lr) 등을 반영해 파라미터 값을 업데이트한다.

✅ 여기서 포인트: optimizer에 파라미터를 잘 등록해야 한다

• optimizer는 “loss를 만드는 데 관여한 텐서”를 자동으로 찾아서 업데이트하지 않는다.
• optimizer는 오직 자기가 들고 있는 파라미터 리스트만 업데이트한다.
• 그래서 보통 이렇게 만든다: optimizer = torch.optim.AdamW(model.parameters(), lr=...)

즉, 학습시키고 싶은 파라미터(nn.Parameter)가 optimizer에 포함되어 있어야 실제 업데이트가 일어난다. (예: 어떤 모듈을 만들고 optimizer에 안 넣으면 그 모듈 파라미터는 grad가 생겨도 업데이트가 안 됨)

4) optimizer.zero_grad() : 다음 step을 위해 grad 초기화

 딥러닝에서 gradient descent로 학습한다는 것은, 항상 “현재 weight 값에서의 loss 기울기”를 이용해 그 위치에서 한 걸음 이동하는 과정을 반복하는 것이다. 이 점이 backpropagation과 optimizer 동작을 이해하는 핵심이다.

이를 수식으로 쓰면, 파라미터 a에 대한 한 step의 업데이트는 다음과 같다.

 

여기서 중요한 것은, 편미분 기호 뒤에 항상 암묵적으로 “a=a_t​에서 계산했다” 는 조건이 붙어 있다는 사실이다.
즉, gradient는 어떤 추상적인 방향이 아니라, 특정 weight 값에서 정의된 국소적인 정보다.

 

PyTorch에서 loss.backward()를 호출하면, 현재 weight 값에서 loss에 대한 각 weight의 gradient가 계산되어 param.grad에 저장된다. 그리고 optimizer.step()은 이 grad 값을 읽어 파라미터를 실제로 업데이트한다. 이 순간 weight 값은 바뀐다.

문제는 그 다음이다. PyTorch는 기본적으로 gradient를 누적(accumulate) 하는 방식을 사용한다. 즉, param.grad는 자동으로 덮어써지지 않고, 새로 계산된 gradient가 이전 값 위에 더해진다. 따라서 optimizer.zero_grad()를 호출하지 않으면, 서로 다른 weight 상태에서 계산된 gradient들이 하나의 grad에 섞이게 된다.

 

이 상황을 간단한 예로 생각해보자.
어떤 파라미터 a가 처음에 a=2일 때 loss에 대한 편미분을 계산했다고 하자.

이 gradient를 이용해 업데이트를 한 뒤, 파라미터가 a=10으로 바뀌었다고 가정하자. 이 새로운 위치에서 다시 loss에 대한 편미분을 계산하면,

가 된다.

정상적인 gradient descent라면, a=10에서의 이동은 오직 g_1​만을 사용해 이루어져야 한다. 즉, 

이 되어야 한다.

하지만 optimizer.zero_grad()를 하지 않으면, PyTorch 내부에서는 업데이트가 다음처럼 이루어진다. 

이 수식이 의미하는 바는 명확하다.
a=2에서 계산된 기울기와 a=10에서 계산된 기울기를 더한 뒤, 그 결과를 a=10에서의 이동 방향으로 사용하고 있는 것이다.

이는 수학적으로 전혀 정당화될 수 없다. 서로 다른 두 지점에서의 접선 기울기를 섞어, 한 지점에서의 하강 방향이라고 주장하는 것과 같기 때문이다. 이 방향은 더 이상 어떤 loss surface의 국소적인 하강 방향도 아니며, 의도하지 않은 업데이트, 즉 학습을 망가뜨리는 “참사”로 이어진다.

 

그래서 optimizer.step()으로 weight가 변한 후에는 꼭 optimizer.zero_grad()를 해주어야 한다. 

 

 

밑의 문장들이 이해가 된다면 진짜 이해한 것이다

 

“같은 weight 상태에서 서로 다른 데이터에 대해 loss를 여러 번 계산하고 loss.backward()를 반복하는 것은, 그 weight 상태에서 더 많은 데이터를 본다는 의미에 불과하다.”

“optimizer.step()이 호출되면 weight 값이 실제로 갱신된다. 이 순간부터 이전 weight 값에서 계산해 두었던 gradient들은 더 이상 현재 weight를 기준으로 한 편미분이 아니므로 의미를 잃는다. 따라서 다음 step에서는 반드시 optimizer.zero_grad()로 gradient를 초기화한 뒤, 새로운 weight 값에서 다시 loss와 gradient를 계산해야 한다.”

 

끝!~