Customizing Autodiff Graphs: Forward and Reverse Mode Automatic Differentiation in Rust
Automatic differentiation is a program transformation method for computing exact derivatives of numerical programs up to floating-point arithmetic error. This document defines a Rust-oriented reference model for forward mode and reverse mode automatic differentiation. The implementation is intentionally scalar and minimal so that operator rules, graph construction, tangent propagation, and adjoint propagation remain visible.
Consider a differentiable function:
\[ f:\mathbb{R}^{n}\rightarrow\mathbb{R}^{m} \]
Its Jacobian is:
\[ J_f(x)= \begin{bmatrix} \frac{\partial f_1}{\partial x_1} & \cdots & \frac{\partial f_1}{\partial x_n} \\ \vdots & \ddots & \vdots \\ \frac{\partial f_m}{\partial x_1} & \cdots & \frac{\partial f_m}{\partial x_n} \end{bmatrix} \]
Forward mode propagates tangents from inputs to outputs. It is efficient when the number of input directions is small. Reverse mode propagates adjoints from outputs to inputs. It is efficient when the output is scalar and the number of input parameters is large. Neural network training commonly uses reverse mode because the objective is usually a scalar loss function:
\[ L:\mathbb{R}^{n}\rightarrow\mathbb{R} \]
For a computational graph with intermediate values \(v_i\), forward mode evaluates:
\[ \dot{v}_j=\sum_i\frac{\partial v_j}{\partial v_i}\dot{v}_i \]
Reverse mode evaluates:
\[ \bar{v}_i=\sum_j\bar{v}_j\frac{\partial v_j}{\partial v_i} \]
The following Rust implementation demonstrates both approaches. The first section uses dual numbers for forward mode. The second section uses a tape structure for reverse mode. This pattern can be used to validate custom differentiable operators before implementing them in a tensor backend.
Reference Implementation
#[derive(Clone, Copy, Debug)]
pub struct Dual {
pub value: f64,
pub tangent: f64,
}
impl Dual {
pub fn variable(value: f64) -> Self {
Self {
value,
tangent: 1.0,
}
}
pub fn constant(value: f64) -> Self {
Self {
value,
tangent: 0.0,
}
}
pub fn sin(self) -> Self {
Self {
value: self.value.sin(),
tangent: self.tangent * self.value.cos(),
}
}
pub fn exp(self) -> Self {
let primal = self.value.exp();
Self {
value: primal,
tangent: self.tangent * primal,
}
}
}
impl std::ops::Add for Dual {
type Output = Dual;
fn add(self, rhs: Dual) -> Dual {
Dual {
value: self.value + rhs.value,
tangent: self.tangent + rhs.tangent,
}
}
}
impl std::ops::Mul for Dual {
type Output = Dual;
fn mul(self, rhs: Dual) -> Dual {
Dual {
value: self.value * rhs.value,
tangent: self.tangent * rhs.value + self.value * rhs.tangent,
}
}
}
#[derive(Clone, Debug)]
enum Op {
Input,
Add(usize, usize),
Mul(usize, usize),
Sin(usize),
Exp(usize),
}
#[derive(Clone, Debug)]
struct Node {
value: f64,
grad: f64,
op: Op,
}
#[derive(Default)]
pub struct Tape {
nodes: Vec<Node>,
}
impl Tape {
pub fn input(&mut self, value: f64) -> usize {
let id = self.nodes.len();
self.nodes.push(Node {
value,
grad: 0.0,
op: Op::Input,
});
id
}
pub fn add(&mut self, a: usize, b: usize) -> usize {
let id = self.nodes.len();
self.nodes.push(Node {
value: self.nodes[a].value + self.nodes[b].value,
grad: 0.0,
op: Op::Add(a, b),
});
id
}
pub fn mul(&mut self, a: usize, b: usize) -> usize {
let id = self.nodes.len();
self.nodes.push(Node {
value: self.nodes[a].value * self.nodes[b].value,
grad: 0.0,
op: Op::Mul(a, b),
});
id
}
pub fn sin(&mut self, a: usize) -> usize {
let id = self.nodes.len();
self.nodes.push(Node {
value: self.nodes[a].value.sin(),
grad: 0.0,
op: Op::Sin(a),
});
id
}
pub fn exp(&mut self, a: usize) -> usize {
let id = self.nodes.len();
self.nodes.push(Node {
value: self.nodes[a].value.exp(),
grad: 0.0,
op: Op::Exp(a),
});
id
}
pub fn backward(&mut self, output: usize) {
self.nodes[output].grad = 1.0;
for i in (0..=output).rev() {
let grad = self.nodes[i].grad;
match self.nodes[i].op.clone() {
Op::Input => {}
Op::Add(a, b) => {
self.nodes[a].grad += grad;
self.nodes[b].grad += grad;
}
Op::Mul(a, b) => {
self.nodes[a].grad += grad * self.nodes[b].value;
self.nodes[b].grad += grad * self.nodes[a].value;
}
Op::Sin(a) => {
self.nodes[a].grad += grad * self.nodes[a].value.cos();
}
Op::Exp(a) => {
self.nodes[a].grad += grad * self.nodes[i].value;
}
}
}
}
pub fn value(&self, id: usize) -> f64 {
self.nodes[id].value
}
pub fn grad(&self, id: usize) -> f64 {
self.nodes[id].grad
}
}
fn main() {
let x = Dual::variable(0.7);
let y = x.sin() * x.exp();
println!("forward value: {}", y.value);
println!("forward derivative: {}", y.tangent);
let mut tape = Tape::default();
let a = tape.input(2.0);
let b = tape.input(3.0);
let product = tape.mul(a, b);
let shifted = tape.add(product, b);
let output = tape.sin(shifted);
tape.backward(output);
println!("reverse value: {}", tape.value(output));
println!("df/da: {}", tape.grad(a));
println!("df/db: {}", tape.grad(b));
}Forward mode stores a primal value and a tangent. A dual value can be represented as:
\[ x \mapsto (x,\dot{x}) \]
Multiplication follows the product rule:
\[ (u,\dot{u})(v,\dot{v})=(uv,\dot{u}v+u\dot{v}) \]
The exponential function follows:
\[ \exp(u,\dot{u})=(e^u,\dot{u}e^u) \]
Reverse mode stores a computation tape. Each node records a primal value, an adjoint accumulator, and the operation that produced it. During the backward pass, the output adjoint is initialized to one:
\[ \bar{L}=1 \]
The tape is then traversed in reverse topological order. Each node distributes its accumulated adjoint to its parents using the local derivative rule. For \(z=uv\), the reverse propagation equations are:
\[ \bar{u} \mathrel{+}= \bar{z}v,\quad \bar{v} \mathrel{+}= \bar{z}u \]
A production-grade autodiff system must also handle broadcasting, tensor shapes, memory reuse, graph pruning, mutation restrictions, and device-specific execution. Nevertheless, the scalar tape model is useful because it exposes the required invariants. Every differentiable operation must define a forward primal computation and a backward adjoint rule. If either rule is incorrect, the resulting gradient will be invalid even when the primal output appears correct.