🤖 Rust와 미분하기 03: 정방향 자동 미분

🔖 Automatic Differentiation Series

1. 💻 Numerical Differentiation
2. 🖊️ Symbolic Differentiation
3. 🤖 Automatic Differentiation

딥러닝을 구현함에 있어서 가장 중요한 요소가 뭘까요? 물론 많은 학문으로 구성된 딥러닝의 특성상 모든 요소들이 다 중요하지만, 그 중에서도 특히 신경써야하는 요소가 있습니다. 이를 찾아내기 위해서 다음의 PyTorch 코드를 살펴봅시다.

net = nn.Sequential(
  nn.Linear(2, 1),
  nn.Sigmoid()
)

# x = ...
# y = ...
# criterion = ...
opt = optim.SGD(net.parameters(), lr=0.01)

opt.zero_grad()
loss = criterion(net(x), y)
loss.backward()
opt.step()

이를 아무런 딥러닝 프레임워크를 쓰지 않고 구현한다고 생각해봅시다. 일단 엄밀하게 같은 구현은 아니지만 Linear와 Sigmoid 함수 자체의 구현은 단순히 행렬곱과 벡터화된 sigmoid 함수를 사용하여 구현할 수 있으므로 net(x)를 만드는 것은 어렵지 않습니다. 다음으로 여기선 criterion이 무엇인지 명시하지는 않았지만 가장 기본적인 MSE를 사용한다면 이 역시 간단합니다. 문제는 opt부터 시작됩니다. SGD를 구현하려면, 어떤 criterion이나 신경망 구조에서도 gradient 즉, 도함수를 구할 수 있어야 합니다.

이를 위해 앞에서 배운 수치미분을 사용할 수 있을 것입니다. 위 경우에는 신경망이 단일 계층으로 되어있고 입력, 출력 차원도 작으므로 크게 문제되진 않겠지만 보통 딥러닝 모델들은 다층으로 되어있으며 입출력 차원도 큰데다 성공적인 훈련을 위해서는 수천, 수만 번 이상 훈련을 해야할 수 있습니다. 그럴 경우, 수치미분으로 인한 사소한 오차가 훈련을 거듭하며 눈덩이처럼 불어날 수 있습니다.

그렇다면 앞에서 배운 기호미분은 어떨까요? 기호미분은 정확한 도함수의 형태를 얻을 수 있기에 부동소수점 오차외에는 별도의 오차가 누적되지 않습니다. 하지만 앞선 글에서도 지적했다시피 기호미분은 한 번 계산할때마다 큰 계산 비용을 요구합니다. 이러한 계산 비용은 수천번 이상의 훈련을 거치면서 기하급수적으로 증가하며, 이는 원하는 시간 내에, 혹은 제한된 메모리 상황에서 훈련을 어렵게 만들 수 있습니다.

앞서 배운 두 가지 미분방법 모두 부적합하다면 대체 어떻게 구현해야할까요? 다행히 과학자들은 딥러닝이 나오기 이전에 이미 이를 위한 도구를 발명해놓았습니다. 바로 자동미분 (Automatic Differentiation) 입니다. 자동미분은 계산방법에 따라 크게 두 가지로 나눌 수 있는데, 입력 변수의 변화에 따라 각 함수에 전파되는 도함수를 계산하는 방식을 정방향 자동미분이라 부르며, 반대로 출력 변수의 변화에 따라 도함수를 계산하는 방식을 역방향 자동미분이라 부릅니다. 각 방식은 서로 다른 장단점을 갖고 있어 경우에 따라 선택해서 사용할 수 있는데, 이번 글에서는 특히 정방향 자동미분에 집중하도록 하겠습니다.

1. Truncated Taylor Series

자동미분이라는 거창한 이름이 붙었지만, 정방향 자동미분의 개념은 대학교 미적분학시간에 배우는 테일러 급수에서 유래합니다. 미분가능한 함수 $f$의 특정 지점 $x_0$에서의 테일러 급수는 다음과 같이 전개됩니다.

$$ f(x) = f(x_0) + f’(x_0)(x - x_0) + \frac{f’’(x_0)}{2!}(x - x_0)^2 + \cdots $$

이때 $x$ 대신 $x$에 충분히 작은 (infinitesimal) 변량 $\epsilon$을 가하고 이를 $x$에 대해서 테일러 급수를 전개하면 다음과 같습니다.

$$ f(x+ \epsilon) = f(x) + f’(x)\epsilon + \mathcal{O}(\epsilon^2) $$

$\epsilon$이 충분히 작다고 가정했으므로 ($|\epsilon| \ll 1$), $\epsilon^2$ 이상의 항들은 무시할 수 있습니다.

$$ f(x + \epsilon) = f(x) + f’(x)\epsilon $$

이렇게 특정 차수까지만 작성하는 테일러 급수를 Truncated Taylor Series라 부릅니다. 꼭 일차항까지만 남겨놓을 필요는 없지만, 대부분의 딥러닝 문제는 도함수까지의 정보만 요구하므로 여기서는 일차항까지만 작성하겠습니다.

이제 $\epsilon$에 계수가 곱해져있는 형태를 생각해봅시다.

$$ f(x + \overline{x} \epsilon) = f(x) + f’(x)\overline{x}\epsilon $$

위 식을 보면 $\epsilon$이 마치 일종의 기저처럼 작용한다는 것을 볼 수 있습니다. $x + \overline{x} \epsilon = x \mathbb{1} + \overline{x} \epsilon$ 으로 보면 $\mathbb{1}$과 $\epsilon$이 기저인 공간에서 수를 나타낸 것처럼 보입니다. 따라서 이를 하나의 수로 취급하고 다음과 같이 간단히 표현해봅시다.

$$ x \mathrel{\rhd} \overline{x} \equiv x + \overline{x}\epsilon $$

이러한 수를 이원수 (Dual number) 라 부르며 무려 1873년에 William Clifford에 의해 창안되었습니다. 이원수는 처음 목적이 자동미분은 아니었지만 최근에 와서는 자동미분을 나타내는 대표적인 수체계로 자리잡았습니다. 이제 이 수를 이용하여 다시 일차항까지의 테일러 급수를 작성해봅시다.

$$ f(x \mathrel{\rhd} \overline{x}) = f(x) \mathrel{\rhd} f’(x)\overline{x} $$

이제 좀 더 간단해졌지만 하나의 문제가 남아있습니다. 함수 $f$의 정의가 명확하지 않다는 것입니다. $x$를 1차원 실수로 가정하고 $x \mathrel{\rhd} \overline{x}$를 일차원 이원수($\mathbb{D}\mathbb{R}$)라 가정하면 좌변에서의 함수 $f$는 $\mathbb{DR} \rightarrow \mathbb{DR}$ 함수이지만 우변에서의 함수 $f$는 $\mathbb{R} \rightarrow \mathbb{R}$ 이므로 정의가 되지 않습니다. 이는 다행히 새로운 연산자인 $\vec{\mathcal{J}}$을 도입하면 쉽게 해결할 수 있습니다.

$$ \vec{\mathcal{J}}f(x \mathrel{\rhd}\overline{x}) = f(x) \mathrel{\rhd} f’(x)\overline{x} $$

이때 연산자 $\vec{\mathcal{J}}$는 다양한 관점에서 해석될 수 있는데, 몇 가지만 아래에 기술하겠습니다.

• $f$라는 실수공간에 펼쳐진 장(field)을 이원수 공간의 장으로 변환한 연산자이므로 미분기하학에서의 Push-forward 개념으로 받아들일 수 있습니다.

• 원래 실수 타입을 받아서 실수 타입을 반환하는 함수 $f$에 작용해서 새로운 함수를 만들어냈으므로 함수형 프로그래밍에서의 lifting 개념으로 받아들일 수 있습니다.

• 실수 함수로 작성된 함수 $f$에 작용하여 새로운 함수 형태를 작성하므로 Source code transformation, Operator overloading 혹은 Multiple dispatch 형태로 이해할 수 있습니다.

이처럼 다양한 방식으로 해석할 수 있지만 위 식의 의미는 하나로 귀결되는데, 바로 “정방향 자동미분"을 표현한 식이라는 것입니다. 단순히 일차항까지 표현한 테일러 급수가 어떻게 미분을 표현하는지 의아하실테니 하나의 예시를 들어봅시다. 예시로는 아주 간단한 함수인 $v = \sin u$를 사용하겠습니다.

1.1. 예시: Sin 함수 자동미분

$$ \begin{aligned} &\begin{aligned} v \mathrel{\rhd} \overline{v} &= \vec{\mathcal{J}}\sin (u \mathrel{\rhd} \overline{u}) \\ &= \sin u \mathrel{\rhd} (\cos u) \overline{u} \end{aligned} \\ \therefore ~ &v = \sin u,\quad \overline{v} = (\cos u) \overline{u} \end{aligned} $$

이로써 $v = \sin u$의 도함수는 $\overline{v} = (\cos u) \overline{u}$이라는 결과를 얻었습니다. 이는 우리가 흔히 사용하는 연쇄 법칙(Chain rule)과 일맥상통하므로 올바른 계산임을 알 수 있습니다. 하지만, 언뜻 보면 $\sin u$의 도함수를 $\cos u$로 알려 준 다음 계산을 한 셈이니 기호미분과 무엇이 다른지 의아할 수 있습니다. 이는 직접 코드를 작성해보면 이해할 수 있습니다.

#[derive(Debug, Copy, Clone)]
struct Dual {
    x: f64,
    dx: f64,
}

trait Sin {
    fn sin_(self) -> Self;
}

impl Sin for f64 {
    fn sin_(self) -> Self {
        self.sin()
    }
}

impl Sin for Dual {
    fn sin_(self) -> Self {
        Dual {
            x: self.x.sin(),
            dx: self.dx * self.x.cos(),
        }
    }
}

fn main() {
    let u = Dual { x: 1.0, dx: 1.0 }; // x at x=1
    let v = u.sin_();
    println!("v: {}, dv: {}", v.x, v.dx);
    // sin(1) = 0.8414..
    // sin'(x) = cos(x) = cos(1) = 0.5403..

    let w = v.sin_();
    println!("v: {}, dv: {}", v.x, v.dx);
    // sin(sin(1)) = 0.7456..
    // (sin(sin(x)))' = cos(sin(1)) * cos(1) = 0.3600..
}

Rust를 이용하여 간단히 이원수 구조체를 구현하고 실수와 이원수 모두에 적용되는 Sin이라는 trait를 작성하여 Method overloading을 이용하여 자동미분을 구현해보았습니다. main함수의 첫번째 예시를 보면 $u = 1 \mathrel{\rhd} 1$로 정의된 것을 볼 수 있는데, 이는 값이 1이고 도함수가 1이라는 의미이므로 $x$를 표현한 것으로 볼 수 있습니다. 이원수의 sin_함수 자체에 도함수 계산이 포함되어있으므로 let v = u.sin_();라는 식 만으로 이미 함숫값과 도함숫값이 계산되고 이는 v.x와 v.dx로 저장됩니다. 이렇게 계산된 도함수에는 부동소수점 오차를 제외한 다른 오차는 존재하지 않으며 계산 비용도 단순히 곱하기 한 번과 코사인 한번을 계산한게 전부입니다.

두 번째 예시를 보면 자동미분의 엄청난 성능을 가늠할 수 있는데, 앞서 계산한 v에 한번 더 sin_을 취하여 w를 정의하였습니다. 즉, $w = \sin(\sin(u))$인데, 함수가 복잡해진것과 관계없이 계산 코드는 동일한 것을 볼 수 있습니다. 그렇게 계산된 결과는 놀랍게도 합성함수의 값과 도함수를 정확히 계산해내는 것을 볼 수 있습니다. 따라서 정방향 자동미분을 이용하면 아주 간단한 도함수 전파 규칙을 명시해놓는 것만으로 복잡하게 합성된 함수를 매우 적은 비용으로 계산할 수 있다는 것을 볼 수 있습니다.

2. 정방향 자동미분 구현

위처럼 자동미분은 함수의 형태와 상관없이 작동했던 수치미분과는 다르게 각 함수의 형태마다 도함수의 전파 규칙을 명시해줘야합니다. 따라서 이번에는 사칙연산을 포함하여 대표적인 함수 몇 가지에 대해 자동미분을 구현해보도록 하겠습니다.

2.1. 사칙연산

$w = u \pm v$로 정의된 $w$에 대해서 아까와 마찬가지로 자동미분을 수행해봅시다.

$$ \begin{aligned} &\begin{aligned} w \mathrel{\rhd} \overline{w} &= u \mathrel{\rhd} \overline{u} + v \mathrel{\rhd} \overline{v} \\ &= (u + v) \mathrel{\rhd} (\overline{u} + \overline{v}) \end{aligned} \\ \therefore ~ &w = u + v,\quad \overline{w} = \overline{u} + \overline{v} \end{aligned} $$

마찬가지로 이번엔 $w = u \times v$에 자동미분을 적용해봅시다.

$$ \begin{aligned} &\begin{aligned} w \mathrel{\rhd} \overline{w} &= (u \mathrel{\rhd} \overline{u}) \times (v \mathrel{\rhd} \overline{v}) \\ &= (u + \overline{u}\epsilon) \times (v + \overline{v}\epsilon) \\ &= uv + (u \overline{v} + \overline{u}v)\epsilon \\ &= (u \times v) \mathrel{\rhd} (u\overline{v} + \overline{u}v) \end{aligned} \\ \therefore ~ &w = u \times v,\quad \overline{w} = u\overline{v} + \overline{u}v \end{aligned} $$

위 두 결과는 각각 우리가 익히 알고있는 미분의 선형성과 라이프니츠 규칙과 정확히 들어맞습니다. 이를 위에서 정의한 Dual 구조체에 적용해봅시다.

use std::ops::{Add, Sub, Mul};

impl Add for Dual {
    type Output = Self;

    fn add(self, rhs: Self) -> Self {
        Self {
            x: self.x + rhs.x,
            dx: self.dx + rhs.dx,
        }
    }
}

impl Sub for Dual {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self {
        Self {
            x: self.x - rhs.x,
            dx: self.dx - rhs.dx,
        }
    }
}

impl Mul for Dual {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self {
        Self {
            x: self.x * rhs.x,
            dx: self.x * rhs.dx + self.dx * rhs.x,
        }
    }
}

fn main() {
    let u = Dual { x: 1f64, dx: 1f64 }; // x at x=1
    let v = Dual { x: 2f64, dx: 4f64 }; // 2x^2 at x=1

    let w = u + v;
    println!("w: {}, dw: {}", w.x, w.dx);
    // w: 3, dw: 5

    let w = u * v;
    println!("w: {}, dw: {}", w.x, w.dx);
    // w = 2x^3 at x=1 = 2
    // dw = 6x^2 at x=1 = 6
}

2.2. 지수, 로그, 다항 함수

다음으로는 지수 로그 함수의 대표적인 함수들인 $y=e^x,\,y=\ln x$와 다항함수인 $y=x^n$에 대해서 자동미분을 수행해봅시다.

1. 지수함수 $$ \begin{aligned} &\begin{aligned} v \mathrel{\rhd} \overline{v} &= \exp(u \mathrel{\rhd} \overline{u}) \\ &= e^u \mathrel{\rhd} e^{u} \overline{u} \end{aligned} \\ \therefore ~ &v = e^u,\quad \overline{v} = e^{u} \overline{u} \end{aligned} $$

2. 로그함수 $$ \begin{aligned} &\begin{aligned} w \mathrel{\rhd} \overline{w} &= \ln(u \mathrel{\rhd} \overline{u}) \\ &= \ln u \mathrel{\rhd} \frac{\overline{u}}{u} \end{aligned} \\ \therefore ~ &w = \ln u,\quad \overline{w} = \frac{\overline{u}}{u} \end{aligned} $$

3. 다항함수 $$ \begin{aligned} &\begin{aligned} w \mathrel{\rhd} \overline{w} &= (u \mathrel{\rhd} \overline{u})^n \\ &= u^n \mathrel{\rhd} n u^{n-1} \overline{u} \end{aligned} \\ \therefore ~ &w = u^n,\quad \overline{w} = n u^{n-1} \overline{u} \end{aligned} $$

앞서 구현한 sin함수와 같이 다른 삼각함수들도 모두 구현하여 하나의 trait을 선언하면 다음과 같습니다.

trait Ops {
    fn exp(self) -> Self;
    fn ln(self) -> Self;
    fn sin(self) -> Self;
    fn cos(self) -> Self;
    fn tan(self) -> Self;
    fn powi(self, n: i32) -> Self;
}

impl Ops for Dual {
    fn exp(self) -> Self {
        Self {
            x: self.x.exp(),
            dx: self.x * self.dx.exp(),
        }
    }

    fn ln(self) -> Self {
        Self {
            x: self.x.ln(),
            dx: self.dx / self.x,
        }
    }

    fn sin(self) -> Self {
        Self {
            x: self.x.sin(),
            dx: self.x.cos() * self.dx,
        }
    }

    fn cos(self) -> Self {
        Self {
            x: self.x.cos(),
            dx: -self.x.sin() * self.dx,
        }
    }

    fn tan(self) -> Self {
        let tan = self.x.tan();
        Self {
            x: tan,
            dx: self.dx * (tan * tan + 1.0),
        }
    }

    fn powi(self, n: i32) -> Self {
        Self {
            x: self.x.powi(n),
            dx: n as f64 * self.x.powi(n - 1) * self.dx,
        }
    }
}

2.3. 시그모이드 함수

위에 있는 함수들에 몇 가지 연산 구현을 더하면 별도로 시그모이드 함수의 전파방법을 구현하지 않더라도 시그모이드 함수에 대한 자동미분을 계산할 수 있습니다.

trait Sigmoid: Sized 
    + Ops
    + Neg<Output=Self>
    + Add<f64, Output=Self> 
where
    f64: Div<Self, Output=Self> {
    fn sigmoid(self) -> Self {
        1f64 / ((-self).exp() + 1f64)
    }
}

impl Sigmoid for Dual {}

fn main() {
    let u = Dual { x: 1.0, dx: 1.0 }; // x at x=1
    let z = u.sigmoid();
    println!("z: {}, dz: {}", z.x, z.dx);
    // sigmoid(1), sigmoid'(x) = sigmoid(1) * (1 - sigmoid(1))
}

주의: 위에 상기한 코드외에도 f64와의 연산이 추가적으로 정의되어야 작동하는 코드입니다. 실제로 이 코드를 실행하고 싶다면 다음 링크를 참고하세요.

github.com/Axect/dual

마치며

지금까지 이원수를 이용한 정방향 자동미분에 대해서 알아보았습니다. 쉽게 설명하기 위해서 굉장히 간단한 경우에 대해서만 다뤘는데, 실제로 잘 작동되는 자동미분 라이브러리를 만들기 위해서는 다음과 같은 사항들을 고려해야합니다.

• 고계 도함수에 대한 자동미분

• 다변수 함수에 대한 자동미분

• 행렬과 벡터 변수를 포함한 함수의 자동미분

이 중 첫 2가지 사항에 대해서는 Rust 수치계산 라이브러리인 Peroxide에 잘 구현되어 있으므로 참고하시면 됩니다. 마지막 사항에 대해서는 역방향 자동미분과 함께 다음 글에서 다루도록 하겠습니다.

참고문헌

• Axect, Peroxide, Github

• Barak A. Pearlmutter, Automatic Differentiation: History and Headroom