backward()는 역전파를 어떻게 할까? (Autograd의 Node)

🤔 Problem

PyTorch의 Tensor는 requires_grad가 True로 되어있을 때, 변화도에 대한 연산의 추적을 가능하게 합니다.

즉, 각 Tensor에 대해서 .grad 속성과 .grad_fn 속성이 생깁니다.

 

.grad는 현재 Tensor에 대해 목적함수가 얼마큼 변했나에 대한 변화도의 값, 즉 미분 값을 담고 있으며,

.grad_fn은 이전 Tensor에 대해서 현재 Tensor를 미분해 줄 때, 어떠한 연산에 대한 미분을 해주어야 하는지 특정 연산에 대한 미분 함수 정보를 담고 있습니다. (값이 아닌 함수 정보임을 유의) 예를 들어, b = a + 2라면, b에는 a에 대해 미분을 할 때, 더하기 연산을 통해 생성이 되었으니 더하기 연산으로 만들어졌다는 것을 알고 있다는 뜻입니다.

 

여기서 제가 궁금한 점이자 이번 포스팅의 주제는 'backward()를 호출했을 때, 도대체 어떻게 .grad를 갱신하는가?'입니다. 너무 당연하게도 이론적으로는 역전파를 수행하도록 하면, .grad에는 변화도에 대한 정보가 들어가는 것이 맞지만, 현재 제가 가지고 있는 정보로는 그건 불가능하다고 판단했기 때문입니다.

 

이렇게 판단한 근거에는 Tensor에서 .grad와 .grad_fn과 같은 속성, 혹은 다른 속성들에서도 현재 Tensor에서는 이전 Tensor에 대한 위치 정보(속성)를 찾을 수 없었기 때문입니다.

 

왜냐하면, 현재 Tensor는 이 위치 정보를 알고 있어야 이 위치에게 어떤 정보든 넘겨줄 수 있기 때문에 역전파를 하기 위해서는 '현재 Tensor는 어떠한 Tensor로부터 나왔는지' 무조건 알고 있어야 합니다.

 

이번 포스팅의 내용은 다소 혼잡스러울 수 있기 때문에 맨 하단의 Summary를 꼭 읽어주시거나 댓글을 남겨주시면 감사하겠습니다 :)

---

🧐 Computational Graph (연산 그래프)

이 문제와 관련해서 제일 밀접한 정보를 가지고 있는 키워드는 Computational Graph입니다. 이에 대한 정보는 PyTorch 한국 사용자 모임에서도 이미 제공을 하고 있습니다.

Forward Propagation을 통해서 구축이 되는 Computational Graph는 동적으로 구축이 됩니다. 동적으로 구축이 되는 이유는 이다음에 어떠한 Forward가 발생할지 모르기 때문에 정적으로 관리되지 않는다는 것입니다.

 

그래서 Computational Graph에서는 현재 Tensor의 이전 Tensor는 무엇인지와 같은 정보를 관리하지는 않습니다.

---

😀 Solution: Autograd의 Node

하지만, 위에서 말한 Computational Graph는 PyTorch에 대한 설명이었을 뿐, 애초에 미분을 작동시키고 있는 Autograd에서는 Computational Graph에 대해 더 깊이 다루고 있을 거라는 생각이 들었고, 이에 기반하여 Autograd의 공식 문서를 찾아보았습니다.

https://autograd.readthedocs.io/en/latest/implementation.html

Implementation — autograd 1.0.0 documentation

 

해당 공식 문서에서는 Autograd가 어떻게 구성이 되어있는지를 나타내며, 크게 Variable, Block, Node로 정리할 수 있었습니다. 다만, 이 3가지에서 관리하고자 하는 것들이 조금씩 다르고, 헷갈릴 수 있는 모호성이 있기 때문에 원하는 정보가 있는 Node에 대해서만 이야기하겠습니다.

 

Node는 Autograd에서 관리하는 Data Structure로 y = f(x)와 같은 하나의 함수를 계산하는 과정에 대해 연산 그래프를 생성하며, 함수 내부의 각 연산에 대한 관계와 Gradient를 Node로 관리합니다. y = g(f(x))도 결국 펼쳐보았을 때, 하나의 함수가 여러 개의 연산으로 엮인 것과 같기 때문에 Node를 하나의 Tensor라 보고 이해해도 무방합니다.

 

Forward가 일어나는 경우, PyTorch가 아닌 PyTorch 내부의 Autograd에서는 Computational Graph를 구축합니다. 이러한 그래프의 정보가 담긴 것이 Autograd의 Node입니다. 마지막 Node인 Loss를 Root로 잡고, 연산의 순방향과 반대인 역방향으로 트리를 형성합니다. 위 그림에서 화살표를 반대로 바꾼 것과 같습니다.

🧐 Node

Node에는 크게 2가지 속성이 있습니다.

 

1. children

2. gradient

 

사실, 이미 childrens라는 속성을 가진 것부터 '아 이전 Tensor에 대한 위치 정보는 Autograd의 Node에서 관리가 되겠구나'라는 걸 짐작할 수 있습니다. children에서 관리하는 것은 현재 Node의 Child Node가 무엇인지, 그리고 Child Node와 어떠한 미분 값(Gradient)을 갖는지에 대해 Jacobian Matrix 형태로 가지고 있습니다.

x=Variable(2) 
y=sin(x)

y.node.childrens=[{'node':x.node, 'jacobian':cos(x.data)}]

그리고, gradient에서는 현재 Node에 대한 root(목적함수)의 미분 값을 관리합니다.

 

공식 문서에서 제공하는 child의 gradient(이전 Tensor의 .grad)가 갱신이 되어가는 과정을 나타내는 코드를 가져왔습니다.

for child in self.childrens:
  node,jacobian=child['node'], child['jacobian']
  new_grad = np.dot(self.gradient, jacobian)
  node.update_gradient(new_grad)

이 코드에서 알 수 있는 건 self(= 현재 Node)에서 child(= 이전 Tensor)에 접근을 하여 현재 Node에 대한 목적 함수의 Gradient와 child Node에 대해 현재 Node를 미분한 값을 곱하여 child Node의 gradient(이전 Tensor의 .grad)가 갱신이 되는 것을 확인할 수 있었습니다. 단순히 곱하는 것으로 미분 값이 구해지는 것은 미분의 연쇄 법칙에 의해서 가능한 것이며 이 말을 수식으로 나타냈을 때는 아래와 같습니다.

 

https://draw-code-boy.tistory.com/517

미분의 연쇄 법칙(Chain Rule)에 대하여

$$ \begin{align} 
\text{self.gradient} = \frac{\delta Loss}{\delta Now} \\ \\
\text{child[`jacobian`]} = \frac{\delta Now}{\delta child} \\ \\
\frac{\delta Loss}{\delta child} = \frac{\delta Loss}{\delta Now}\cdot\frac{\delta Now}{\delta child}
\end{align} $$

 

물론, Autograd에서도 Computational Graph는 어떠한 다른 Forward가 발생할지 모르기 때문에 동적으로 관리가 됩니다.

---

✅ Summary

이번 내용이 복잡했던 만큼 제가 찾아보고자 했던 정보에 대해서 간단하게 정리를 하자면, PyTorch에서 Back Propagation이 일어날 때 현재 Tensor에서 이전 Tensor에게 Gradient와 같은 정보들을 넘길 때, 이전 Tensor의 위치를  어떻게 알고있는가에 대한 의문을 품었었습니다.

 

이에 대해서는 Autograd의 Computational Graph를 통해서 현재 Node의 Child Node 정보로 관리가 되고 있다는 사실을 공식 문서를 통해 확인했습니다.