28.34 텐서 연산의 자동 미분(Automatic Differentiation)과 계산 그래프

1. 자동 미분의 개념적 위치

자동 미분(Automatic Differentiation, AD)은 컴퓨터로 표현 가능한 임의의 수치 함수에 대하여, 사용자가 명시적으로 도함수 식을 유도하지 않고도 정확한 도함수의 값을 기계적·체계적으로 계산하는 일군의 알고리즘이다. 이는 기호 미분(Symbolic Differentiation)이 표현식의 폭발을 일으키고, 수치 미분(Numerical Differentiation)이 절단 오차와 반올림 오차의 균형 문제로 인해 정확도가 낮다는 한계를 동시에 극복한다. 자동 미분은 어떠한 근사 오차도 도입하지 않으며, 부동소수점 연산의 본질적 정밀도 한계 내에서 정확한 값을 산출한다는 점에서 본질적으로 정밀(Exact)하다. 그리피스벤크(A. Griewank)와 발터(A. Walther)의 표준 교과서 “Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation“는 이러한 관점을 체계적으로 정립한 대표적 저작이다.

딥러닝의 학습은 손실 함수를 모든 매개변수에 대해 미분하여 경사 하강을 수행하는 절차로 환원되며, 현대 딥러닝 프레임워크의 핵심 동력은 텐서 연산에 대한 효율적이고 정확한 자동 미분 기능에 있다. 자동 미분이 없다면 수십억 개에 달하는 매개변수를 가진 신경망의 도함수를 손으로 유도하는 것은 사실상 불가능하다.

2. 계산 그래프의 정의

자동 미분의 모든 형태는 함수의 계산을 계산 그래프(Computational Graph)라는 자료 구조로 표현하는 것에서 출발한다. 계산 그래프는 방향성 비순환 그래프(Directed Acyclic Graph, DAG) $G = (V, E)$ 로 정의되며, 각 정점은 텐서 변수 또는 원시 연산(Primitive Operation)에 해당하고, 각 변은 변수 사이의 함수적 의존 관계를 나타낸다. 정점은 일반적으로 세 가지 종류로 구분된다. 첫째는 입력 정점으로, 외부에서 주어지는 텐서 자료나 학습 가능한 매개변수에 해당한다. 둘째는 중간 정점으로, 어떤 원시 연산의 출력을 표현한다. 셋째는 출력 정점으로, 최종적으로 도함수를 계산하고자 하는 스칼라 값(예: 손실 함수)을 표현한다.

계산 그래프의 각 변은 자식 정점이 부모 정점에 함수적으로 의존한다는 사실을 부호화하며, 그래프 전체는 입력으로부터 출력에 이르는 전체 계산을 일련의 원시 연산의 합성(Composition)으로 분해한 형태가 된다. 자동 미분 알고리즘은 이 그래프 위에서 연쇄 법칙(Chain Rule)을 체계적으로 적용함으로써 도함수를 계산한다.

3. 전향 모드와 역향 모드

자동 미분은 연쇄 법칙을 적용하는 방향에 따라 전향 모드(Forward Mode)와 역향 모드(Reverse Mode)의 두 가지 형태로 구분된다.

3.1 전향 모드

함수 $f: \mathbb{R}^{n} \to \mathbb{R}^{m}$ 에 대하여, 입력 변수에 대한 미소 변화를 출력 방향으로 전파하는 방식이다. 임의의 방향 벡터 $\mathbf{v} \in \mathbb{R}^{n}$ 이 주어졌을 때, 전향 모드는 야코비안-벡터 곱(Jacobian-Vector Product, JVP) $J_f(\mathbf{x}) \, \mathbf{v}$ 를 계산한다. 형식적으로는 각 중간 변수 $u_i$ 에 대해 그 값과 함께 입력 방향 $\mathbf{v}$ 에 대한 방향 도함수 $\dot{u}_i$ 를 동시에 추적하며, 원시 연산의 미분 규칙에 따라 두 양을 함께 갱신한다. 한 번의 전향 모드 실행은 한 개의 입력 방향에 대한 모든 출력의 도함수를 산출한다. 따라서 전향 모드는 입력 차원이 작고 출력 차원이 큰 함수에 효율적이다.

3.2 역향 모드

이와 반대로, 출력 측의 미소 변화를 입력 방향으로 전파하는 방식이다. 임의의 보조 벡터 $\mathbf{w} \in \mathbb{R}^{m}$ 이 주어졌을 때, 역향 모드는 벡터-야코비안 곱(Vector-Jacobian Product, VJP) $\mathbf{w}^{\top} J_f(\mathbf{x})$ 를 계산한다. 손실 함수가 스칼라 출력을 가지는 경우 $\mathbf{w}$ 는 단일 1로 환원되며, 이때 한 번의 역향 모드 실행만으로 모든 입력에 대한 손실의 경사도가 동시에 산출된다. 따라서 역향 모드는 입력 차원이 매우 크고 출력 차원이 작은 함수, 즉 거대한 매개변수 공간 위에서 정의된 스칼라 손실 함수의 경사를 구하는 딥러닝의 표준 상황에 가장 적합하다. 일반적으로 딥러닝에서 “자동 미분“이라는 용어는 사실상 역향 모드 자동 미분을 의미하며, 이는 통상적인 호칭인 역전파(Backpropagation)와 동일한 알고리즘이다.

4. 역향 모드의 형식적 절차

함수 $L = f(\mathbf{x})$ 가 일련의 중간 변수 $u_1, u_2, \ldots, u_K = L$ 로 구성되어 있다고 가정한다. 역향 모드는 다음의 두 단계로 진행된다.

첫째, 전향 단계(Forward Pass)에서 입력으로부터 출력 방향으로 그래프를 따라가며 모든 중간 변수의 값을 차례로 계산하고, 각 원시 연산의 입력값을 메모리에 저장한다. 이는 후행 단계에서 국소 도함수(Local Derivative)를 계산하기 위한 준비 단계이다.

둘째, 후행 단계(Backward Pass)에서 출력으로부터 입력 방향으로 그래프를 거꾸로 따라가며, 각 중간 변수에 대한 출력의 도함수, 즉 보조량(Adjoint) $\bar{u}_i = \partial L / \partial u_i$ 를 차례로 계산한다. 보조량은 다음의 점화식에 따라 갱신된다.

$\bar{u}_i = \sum_{j \,:\, u_i \to u_j} \bar{u}_j \cdot \frac{\partial u_j}{\partial u_i}$

여기서 합산은 $u_i$ 가 직접적으로 입력으로 사용된 모든 자식 정점 $u_j$ 에 대해 이루어진다. 출력 정점 $L$ 에 대해서는 $\bar{L} = 1$ 로 초기화되며, 이 식을 입력 정점에 도달할 때까지 반복적으로 적용함으로써 모든 중간 변수와 입력 변수에 대한 도함수를 동시에 얻는다.

5. 텐서 연산에서의 자동 미분

딥러닝의 자동 미분이 통상적인 스칼라 자동 미분과 본질적으로 구별되는 점은, 모든 변수가 텐서이며 모든 원시 연산이 텐서 연산이라는 사실이다. 따라서 야코비안 자체를 명시적으로 구성하지 않고, 야코비안과 벡터의 곱만을 효율적으로 계산하는 형태로 알고리즘이 구현된다. 즉, 각 원시 텐서 연산에는 그에 대응하는 벡터-야코비안 곱(VJP) 규칙이 정의되어 있으며, 후행 단계는 이 VJP 규칙들을 그래프의 위상적 역순으로 합성하는 절차에 해당한다.

예를 들어 행렬 곱셈 $\mathcal{Y} = \mathcal{X} W$ 에 대한 VJP는 다음과 같이 도출된다. 손실 $L$ 의 출력에 대한 보조량 $\bar{\mathcal{Y}} = \partial L / \partial \mathcal{Y}$ 가 주어졌을 때, 입력에 대한 보조량은 $\bar{\mathcal{X}} = \bar{\mathcal{Y}} W^{\top}$ 이고, 가중치에 대한 보조량은 $\bar{W} = \mathcal{X}^{\top} \bar{\mathcal{Y}}$ 이다. 어떠한 야코비안 행렬도 명시적으로 형성되지 않으며, 모든 계산은 동일한 텐서 연산의 변형으로 환원된다.

6. 정적 그래프와 동적 그래프

계산 그래프의 구축 시점에 따라 자동 미분 시스템은 정적 그래프(Static Graph)와 동적 그래프(Dynamic Graph)의 두 가지 패러다임으로 분류된다.

6.1 정적 그래프

전체 계산 그래프를 코드 실행 이전에 한 번에 정의하고, 이후 동일한 그래프에 대해 반복적으로 데이터를 흘려보내며 학습을 수행하는 방식이다. TensorFlow 1.x의 tf.Graph 및 tf.Session 모형이 대표적 사례이다. 정적 그래프는 그래프 전체에 대한 전역적 최적화(연산 융합, 메모리 재사용, 정수형 양자화 등)와 그래프의 분산 실행, 효율적인 직렬화에 유리하다. 그러나 제어 흐름이 데이터에 의존하는 계산을 표현하기가 번거롭고, 디버깅 시 일반적인 명령형 디버거의 사용이 제한된다는 단점을 가진다.

6.2 동적 그래프

각 순전파 단계에서 그래프가 즉석에서 구성되고, 후행 단계가 끝난 직후 그래프가 폐기되는 방식이다. PyTorch의 autograd 모형이 대표적이며, “Define-by-Run” 또는 “Eager Execution“이라 불린다. 동적 그래프는 일반적인 파이썬 제어 흐름을 그대로 사용할 수 있고, 표준 디버거와 자연스럽게 결합되며, 데이터 의존적인 그래프 구조를 매우 간결하게 표현할 수 있다. 현대 프레임워크는 동적 그래프의 사용 편의성과 정적 그래프의 최적화 이점을 결합하기 위하여, 동적으로 추적된 계산을 사후에 정적 그래프로 변환하는 추적 컴파일(Tracing Compilation) 또는 즉시 실행 컴파일(Just-In-Time Compilation) 기법을 도입하고 있다. PyTorch의 torch.compile, JAX의 jit, TensorFlow 2.x의 tf.function이 이러한 융합의 대표적 사례이다.

7. 메모리 비용과 체크포인트 기법

역향 모드 자동 미분의 본질적 특징 가운데 하나는, 후행 단계에서 사용될 중간 활성값을 모두 메모리에 저장해야 한다는 점이다. 이는 깊이가 매우 큰 신경망에서 메모리 사용량이 깊이에 비례하여 선형적으로 증가하는 결과를 낳는다. 이 문제를 완화하기 위하여 도입된 기법이 경사 체크포인팅(Gradient Checkpointing)이다. 체크포인팅은 중간 활성값의 일부만을 저장하고, 후행 단계에서 필요한 시점에 저장되지 않은 부분을 다시 계산하여 메모리와 계산 시간 사이의 절충을 수행한다. 첸(T. Chen) 등의 2016년 논문 “Training Deep Nets with Sublinear Memory Cost“는 이 기법을 통해 깊이 $L$ 의 신경망에서 메모리 사용량을 $\mathcal{O}(\sqrt{L})$ 로 축소할 수 있음을 보였다.

8. 고차 미분과 확장

자동 미분은 도함수에만 국한되지 않고, 도함수에 대한 도함수, 즉 고차 미분(Higher-Order Derivative)에도 동일하게 적용된다. 일반적으로 1차 도함수의 계산 그래프 자체에 다시 자동 미분을 적용함으로써 2차 도함수를 얻을 수 있으며, 이는 헤시안-벡터 곱(Hessian-Vector Product)과 같은 양을 야코비안 행렬을 명시적으로 구성하지 않고도 계산할 수 있게 한다. 펄머터(B. Pearlmutter)의 1994년 논문 “Fast Exact Multiplication by the Hessian“이 이러한 기법의 고전적 분석이며, 현대 프레임워크는 이러한 고차 미분을 자연스럽게 지원한다.

9. 자동 미분과 텐서 추상화의 결합

결국 자동 미분은 단순한 미분 알고리즘이 아니라, 텐서 연산이라는 추상화 위에 정확한 도함수 계산 능력을 부여하는 메타-수준의 기능이다. 이로써 사용자는 신경망을 단순한 텐서 연산의 합성으로 기술하기만 하면 되고, 학습에 필요한 모든 도함수 계산은 프레임워크가 계산 그래프 위에서 자동으로 처리한다. 이러한 분리는 모델 설계와 최적화 알고리즘의 분리(Separation of Concerns)를 가능하게 하며, 딥러닝이 폭발적으로 발전한 가장 중요한 인프라적 기반 가운데 하나이다.