rustytorch/docs/book/src/architecture/autograd-engine.md

Autograd Engine

How automatic differentiation works in RustyTorch++.

Overview

The autograd engine implements reverse-mode automatic differentiation (backpropagation) using a tape-based approach.

Key Concepts

Computation Graph

Every tensor operation builds a computation graph:

     ┌───────────────┐
     │   Input X     │  (leaf, requires_grad=true)
     └───────┬───────┘
             │
     ┌───────▼───────┐
     │   MatMul(W)   │  W is another leaf
     └───────┬───────┘
             │
     ┌───────▼───────┐
     │     ReLU      │
     └───────┬───────┘
             │
     ┌───────▼───────┐
     │     Loss      │  (output)
     └───────────────┘

Gradient Tape

Operations are recorded on a tape:

use rtx_autograd::{AutogradTape, TapeNode};

// Tape records operations during forward pass
let tape = AutogradTape::new();

// Each operation adds a node
tape.record_operation(
    &[input_id, weight_id],  // inputs
    output_id,                // output
    Box::new(MatMulBackward { /* saved tensors */ }),
);

Using Variables

Basic Gradient Computation

use rtx_autograd::{Variable, backward};
use rtx_tensor::{Tensor, Device};

let device = Device::cpu();

// Create tracked variables
let x = Variable::from_tensor(Tensor::from_slice(&[2.0, 3.0], [2], &device)?);
let w = Variable::from_tensor(Tensor::from_slice(&[1.0, 2.0], [2], &device)?);

// Forward pass builds graph
let y = x.multiply(&w)?;  // [2.0, 6.0]
let z = y.sum()?;         // 8.0

// Backward pass computes gradients
z.backward()?;

// Access gradients
println!("dx: {:?}", x.grad());  // [1.0, 2.0] (dy/dx = w)
println!("dw: {:?}", w.grad());  // [2.0, 3.0] (dy/dw = x)

Gradient Context Managers

use rtx_autograd::{no_grad, enable_grad, is_grad_enabled};

// Check current state
println!("Grad enabled: {}", is_grad_enabled());

// Disable gradients (inference mode)
let result = no_grad(|| {
    let a = Tensor::randn([100, 100], &device)?;
    let b = Tensor::randn([100, 100], &device)?;
    a.matmul(&b)  // No graph built
});

// Re-enable within nested scope
no_grad(|| {
    enable_grad(|| {
        // Gradients tracked here
    });
    // Gradients not tracked here
});

Detaching from Graph

let x = Variable::from_tensor(tensor);
let y = x.multiply(&w)?;

// Detach creates a new variable without graph connection
let y_detached = y.detach();
assert!(!y_detached.requires_grad());

// Useful for stop-gradient operations
let target = predictions.detach();  // Don't backprop through target

Backward Functions

Each operation has a corresponding backward function:

use rtx_autograd::{BackwardFunction, BackwardContext};

/// Backward function for matrix multiplication: C = A @ B
pub struct MatMulBackward {
    a: Tensor,  // Saved input A
    b: Tensor,  // Saved input B
}

impl BackwardFunction for MatMulBackward {
    fn backward(&self, grad_output: &Tensor) -> Vec<Option<Tensor>> {
        // dL/dA = dL/dC @ B^T
        let grad_a = grad_output.matmul(&self.b.transpose(0, 1)?)?;

        // dL/dB = A^T @ dL/dC
        let grad_b = self.a.transpose(0, 1)?.matmul(grad_output)?;

        vec![Some(grad_a), Some(grad_b)]
    }
}

Built-in Backward Functions

Operation	Backward	Notes
Add	[1, 1]	Gradients pass through unchanged
Mul	[y, x]	Swap inputs
MatMul	[grad @ B^T, A^T @ grad]	Standard matrix calculus
Sum	expand	Broadcast gradient
ReLU	grad * (x > 0)	Zero where input was negative
Sigmoid	grad * σ * (1 - σ)	Uses saved output
Softmax	Jacobian-vector product	Complex formula

Gradient Checkpointing

Trade memory for compute by recomputing activations:

use rtx_autograd::{CheckpointManager, CheckpointStrategy};

// Create checkpoint manager
let mut checkpoint = CheckpointManager::new(CheckpointStrategy::SqrtN);

// Checkpoint expensive operations
let hidden = checkpoint.checkpoint(|| {
    model.layer1.forward(&input)?
        .relu()?
        .layer2.forward()
})?;

// During backward, activations are recomputed
hidden.backward()?;  // Recomputes forward pass for gradients

Checkpointing Strategies

Strategy	Memory	Compute	Use Case
None	O(n)	O(1)	Small models
SqrtN	O(√n)	O(√n)	Large transformers
Every(k)	O(n/k)	O(k)	Custom balance
Selective	Varies	Varies	Manual selection

Higher-Order Derivatives

Hessian Computation

use rtx_autograd::{HessianComputer};

let f = |x: &Variable| -> Variable {
    x.pow(2)?.sum()?  // f(x) = sum(x^2)
};

let x = Variable::from_tensor(Tensor::from_slice(&[1.0, 2.0, 3.0], [3], &device)?);

// Compute Hessian (second derivative matrix)
let hessian = HessianComputer::compute(&f, &x)?;
// For sum(x^2), Hessian is 2*I (diagonal of 2s)

Jacobian Computation

use rtx_autograd::JacobianComputer;

let f = |x: &Variable| -> Variable {
    // f: R^3 -> R^2
    let a = x.narrow(0, 0, 2)?;
    let b = x.narrow(0, 1, 2)?;
    a.add(&b)?
};

let x = Variable::from_tensor(Tensor::randn([3], &device)?);
let jacobian = JacobianComputer::compute(&f, &x)?;  // [2, 3] matrix

Numerical Gradient Checking

Verify gradient implementations with finite differences:

use rtx_autograd::GradientChecker;

let checker = GradientChecker::new(1e-5);  // epsilon for finite diff

let f = |x: &Tensor| {
    x.pow(3)?.sum()  // f(x) = sum(x^3)
};

let x = Tensor::randn([10], &device)?;
let analytical_grad = compute_analytical_grad(&f, &x)?;

// Check matches numerical gradient
let is_correct = checker.check(&f, &x, &analytical_grad)?;
assert!(is_correct, "Gradient implementation is incorrect");

Graph Caching

Cache computation graphs for repeated forward passes:

use rtx_autograd::{init_graph_cache, cache_graph, get_cached_graph};

// Initialize graph cache
init_graph_cache(1024);  // Cache up to 1024 graphs

// First forward pass - builds and caches graph
let output = model.forward(&input)?;

// Subsequent passes with same shapes - reuse graph
let cached = get_cached_graph(&input_shape)?;
if let Some(graph) = cached {
    // Execute cached graph
    output = graph.execute(&input)?;
}

Best Practices

1. Use no_grad for Inference

// Wrong - builds unnecessary graph
let output = model.forward(&input)?;

// Right - skips graph construction
let output = no_grad(|| model.forward(&input))?;

2. Zero Gradients Before Each Step

for epoch in 0..num_epochs {
    for (x, y) in &dataloader {
        // Clear accumulated gradients
        optimizer.zero_grad();

        let output = model.forward(&x)?;
        let loss = loss_fn(&output, &y)?;

        loss.backward()?;
        optimizer.step()?;
    }
}

3. Detach When Needed

// Stop gradient flow to targets
let target = teacher_model.forward(&x)?.detach();
let student_out = student_model.forward(&x)?;
let loss = mse_loss(&student_out, &target)?;

4. Use Checkpointing for Large Models

// Memory-efficient transformer
let attention = checkpoint.checkpoint(|| {
    self.attention.forward(&hidden)?
})?;

Next Steps

• GPU Runtime - Device and memory management
• Performance - Optimization
• Troubleshooting - Memory issues