Rust의 std::autodiff 모듈을 사용하면 미분 가능 프로그래밍(differentiable programming)을 할 수 있습니다:
#![feature(autodiff)]
use std::autodiff::*;
// f(x) = x * x, f'(x) = 2.0 * x
// bar therefore returns (x * x, 2.0 * x)
#[autodiff_reverse(bar, Active, Active)]
fn foo(x: f32) -> f32 { x * x }
fn main() {
assert_eq!(bar(3.0, 1.0), (9.0, 6.0));
assert_eq!(bar(4.0, 1.0), (16.0, 8.0));
}std::autodiff 모듈에 대한 자세한 문서는 std::autodiff에서 확인할 수 있습니다.
미분 가능 프로그래밍은 수치 계산, 고체 역학, 계산 화학, 유체 역학 등 다양한 분야에서 사용되며, 역전파를 통한 신경망 학습, ODE 솔버, 미분 가능 렌더링, 양자 컴퓨팅, 기후 시뮬레이션에도 사용됩니다.
std::autodiff는 현재 자동 미분을 위한 LLVM 기반 도구인 Enzyme을 기반으로 합니다. 컴파일러 기반 autodiff를 사용하는 주된 이유는 세 가지입니다.
- 사용성: 현재의 autodiff 크레이트는 일반적인 Rust 프로그램을 지원하지 않습니다. 이들 크레이트는 사용자 정의 DSL을 강제하거나, (슬라이스나 배열 등 대신) 라이브러리가 제공하는 타입의 사용을 요구하거나, 스칼라 함수로만 제한됩니다. 컴파일러 기반 autodiff를 사용하면 배열, 슬라이스, 사용자 정의 구조체와 열거형, 제어 흐름 등을 포함한 일반적인 Rust 코드를 작성할 수 있습니다.
- 성능: 기존 Rust autodiff 접근 방식은 대부분 연산마다 일정한 오버헤드가 있습니다. 이 오버헤드는 큰 텐서에 대해 비용이 큰 연산을 적은 횟수로 수행하는 ML 애플리케이션에서는 쉽게 상쇄될 수 있습니다. 그러나 HPC나 과학 계산 분야의 애플리케이션에서는 받아들일 수 없는 경우가 많습니다. 컴파일러 기반 autodiff는 (최적화된) LLVM IR을 대상으로 동작하므로 이러한 경우에 훨씬 더 나은 성능을 낼 수 있습니다.
- 기능: 이처럼 낮은 수준에서 동작하고 다른 LLVM 기반 언어와 구현을 공유하므로, Enzyme 프로젝트에서 이미 이루어진 방대한 작업을 활용할 수 있습니다. 예를 들어, MPI 루틴을 호출하는 Rust 코드나 CuBLAS 같은 라이브러리를 포함한 GPU 코드를 지원할 수 있습니다.