Backward Pass

The backward pass computes gradients by traversing the computational graph in reverse order, applying the chain rule at each operation. This chapter explains the mechanics of gradient propagation in Entrenar.

The Chain Rule

The foundation of backpropagation is the multivariate chain rule:

Given: z = f(y) and y = g(x)
Then:  dz/dx = dz/dy * dy/dx

For neural networks with many layers:

Loss = f_n(f_{n-1}(...f_2(f_1(x))))

dLoss/dx = dLoss/df_n * df_n/df_{n-1} * ... * df_2/df_1 * df_1/dx

Entrenar automates this chain rule application.

Backward Pass Algorithm

High-Level Steps

Seed the gradient: Set output gradient to 1.0
Traverse in reverse: Process tape entries from end to start
Apply local gradients: Each operation computes input gradients from output gradient
Accumulate gradients: Sum contributions when tensors have multiple consumers

Pseudocode

def backward(output_tensor):
    # Step 1: Seed gradient
    output_tensor.grad = 1.0

    # Step 2: Get tape entries
    tape = get_global_tape()

    # Step 3: Reverse traversal
    for entry in reversed(tape):
        # Get output gradient (already computed)
        grad_output = entry.output.grad

        # Step 4: Compute input gradients (chain rule)
        grad_inputs = entry.operation.backward(grad_output)

        # Step 5: Accumulate into input tensors
        for (input_tensor, grad_input) in zip(entry.inputs, grad_inputs):
            input_tensor.grad += grad_input  # Accumulation!

Operation-Specific Backward Rules

Each operation implements a backward method that computes input gradients from output gradients.

Addition: z = x + y

Forward: z_i = x_i + y_i

Backward:

∂z/∂x = 1  (gradient passes through unchanged)
∂z/∂y = 1

Therefore:
∂Loss/∂x = ∂Loss/∂z * 1 = ∂Loss/∂z
∂Loss/∂y = ∂Loss/∂z * 1 = ∂Loss/∂z

Implementation:

#![allow(unused)]
fn main() {
fn add_backward(grad_output: &[f32], x: &Tensor, y: &Tensor) {
    // Gradient flows equally to both inputs
    x.accumulate_grad(grad_output);  // dx = dz
    y.accumulate_grad(grad_output);  // dy = dz
}
}

Multiplication: z = x * y

Forward: z_i = x_i * y_i

Backward:

∂z/∂x = y  (gradient scaled by other input)
∂z/∂y = x

Therefore:
∂Loss/∂x = ∂Loss/∂z * y
∂Loss/∂y = ∂Loss/∂z * x

Implementation:

#![allow(unused)]
fn main() {
fn mul_backward(grad_output: &[f32], x: &Tensor, y: &Tensor) {
    // Gradient to x scaled by y's value
    let grad_x: Vec<f32> = grad_output.iter()
        .zip(y.data().iter())
        .map(|(g, y_val)| g * y_val)
        .collect();
    x.accumulate_grad(&grad_x);

    // Gradient to y scaled by x's value
    let grad_y: Vec<f32> = grad_output.iter()
        .zip(x.data().iter())
        .map(|(g, x_val)| g * x_val)
        .collect();
    y.accumulate_grad(&grad_y);
}
}

Matrix Multiplication: C = A @ B

Forward: C = A @ B (dimensions: C[m,n] = A[m,k] @ B[k,n])

Backward:

∂Loss/∂A = ∂Loss/∂C @ B^T
∂Loss/∂B = A^T @ ∂Loss/∂C

Derivation (element-wise):

C[i,j] = Σ_k A[i,k] * B[k,j]

∂C[i,j]/∂A[i,k] = B[k,j]  => ∂Loss/∂A[i,k] = Σ_j ∂Loss/∂C[i,j] * B[k,j]
                                             = (∂Loss/∂C @ B^T)[i,k]

∂C[i,j]/∂B[k,j] = A[i,k]  => ∂Loss/∂B[k,j] = Σ_i ∂Loss/∂C[i,j] * A[i,k]
                                             = (A^T @ ∂Loss/∂C)[k,j]

Implementation:

#![allow(unused)]
fn main() {
fn matmul_backward(
    grad_output: &Tensor,  // dC
    a: &Tensor,            // A
    b: &Tensor,            // B
    m: usize,              // rows of A
    k: usize,              // cols of A = rows of B
    n: usize,              // cols of B
) {
    // dA = dC @ B^T
    let b_transpose = transpose(b, k, n);
    let grad_a = matmul(grad_output, &b_transpose, m, n, k);
    a.accumulate_grad(grad_a.data());

    // dB = A^T @ dC
    let a_transpose = transpose(a, m, k);
    let grad_b = matmul(&a_transpose, grad_output, k, m, n);
    b.accumulate_grad(grad_b.data());
}
}

ReLU: y = max(0, x)

Forward: y_i = max(0, x_i)

Backward:

∂y/∂x = {1 if x > 0, 0 otherwise}

Therefore:
∂Loss/∂x_i = ∂Loss/∂y_i * (x_i > 0 ? 1 : 0)

Implementation:

#![allow(unused)]
fn main() {
fn relu_backward(grad_output: &[f32], x: &Tensor) {
    let grad_x: Vec<f32> = grad_output.iter()
        .zip(x.data().iter())
        .map(|(g, &x_val)| {
            if x_val > 0.0 {
                *g  // Gradient passes through
            } else {
                0.0  // Gradient blocked
            }
        })
        .collect();

    x.accumulate_grad(&grad_x);
}
}

GELU: y = x * Φ(x)

Forward: y = x * Φ(x) where Φ is the Gaussian CDF

Backward (using product rule):

∂y/∂x = Φ(x) + x * φ(x)

where φ(x) = (1/√(2π)) * exp(-x²/2) is the Gaussian PDF

Implementation:

#![allow(unused)]
fn main() {
fn gelu_backward(grad_output: &[f32], x: &Tensor) {
    const SQRT_2_PI: f32 = 2.5066282746;  // √(2π)

    let grad_x: Vec<f32> = grad_output.iter()
        .zip(x.data().iter())
        .map(|(g, &x_val)| {
            let phi = gaussian_cdf(x_val);          // Φ(x)
            let phi_prime = (-0.5 * x_val.powi(2)).exp() / SQRT_2_PI;  // φ(x)
            let local_grad = phi + x_val * phi_prime;

            g * local_grad
        })
        .collect();

    x.accumulate_grad(&grad_x);
}
}

Layer Normalization

Forward:

y = (x - μ) / σ

where:
  μ = mean(x)
  σ = √(variance(x) + ε)

Backward (complex chain rule):

∂Loss/∂x_i = (1/σ) * [∂Loss/∂y_i - (1/n)Σ_j ∂Loss/∂y_j - (1/n)y_i Σ_j(∂Loss/∂y_j * y_j)]

Implementation:

#![allow(unused)]
fn main() {
fn layernorm_backward(
    grad_output: &[f32],
    x: &Tensor,
    normalized: &[f32],  // y values from forward pass
    mean: f32,
    variance: f32,
) {
    let n = grad_output.len() as f32;
    let std_inv = 1.0 / (variance + 1e-5).sqrt();

    // Compute sum terms
    let sum_grad: f32 = grad_output.iter().sum();
    let sum_grad_y: f32 = grad_output.iter()
        .zip(normalized.iter())
        .map(|(g, y)| g * y)
        .sum();

    // Compute gradient for each element
    let grad_x: Vec<f32> = grad_output.iter()
        .zip(normalized.iter())
        .map(|(g, y)| {
            std_inv * (g - sum_grad / n - y * sum_grad_y / n)
        })
        .collect();

    x.accumulate_grad(&grad_x);
}
}

Gradient Accumulation

When a tensor is used multiple times, gradients accumulate:

Example: z = x + x

#![allow(unused)]
fn main() {
let x = Tensor::from_vec(vec![2.0], true);
let z = &x + &x;  // z = 2x

backward(&z);

println!("dz/dx = {}", x.grad()[0]);  // 2.0 ✅
}

Why 2.0?

Graph:
    x ─┬─> Add -> z
       └─>

Backward:
  From first input:  dx = dz * 1 = 1.0
  From second input: dx = dz * 1 = 1.0
  Total:             dx = 1.0 + 1.0 = 2.0 ✅

Implementation:

#![allow(unused)]
fn main() {
// Always use += for gradient accumulation
x.grad_mut()[i] += gradient_contribution;
}

Complex Example

#![allow(unused)]
fn main() {
let x = Tensor::from_vec(vec![3.0], true);
let y = Tensor::from_vec(vec![4.0], true);

let a = &x + &y;   // a = x + y = 7
let b = &x * &y;   // b = x * y = 12
let c = &a + &b;   // c = a + b = 19

backward(&c);
}

Gradient computation:

Tape (forward order):
  [0] Add(x, y) -> a
  [1] Mul(x, y) -> b
  [2] Add(a, b) -> c

Backward (reverse order):
  [2] Add: da = dc = 1.0, db = dc = 1.0
  [1] Mul: dx += db * y = 1.0 * 4 = 4.0
           dy += db * x = 1.0 * 3 = 3.0
  [0] Add: dx += da = 1.0
           dy += da = 1.0

Final gradients:
  dx = 4.0 + 1.0 = 5.0  ✅ (= y + 1)
  dy = 3.0 + 1.0 = 4.0  ✅ (= x + 1)

Manual verification:

c = (x + y) + (x * y) = x + y + xy
dc/dx = 1 + y = 1 + 4 = 5.0 ✅
dc/dy = 1 + x = 1 + 3 = 4.0 ✅

Handling Non-Differentiable Points

Some operations have non-differentiable points where we use subgradients.

ReLU at x=0

ReLU(x) = max(0, x)

Derivative:
  d/dx ReLU(x) = {1 if x > 0, 0 if x < 0, ??? if x = 0}

Solution: Use subgradient convention:

#![allow(unused)]
fn main() {
if x_val > 0.0 {
    1.0
} else {
    0.0  // Subgradient at x=0 (could also use 1.0 or 0.5)
}
}

In practice: Exact x=0 is rare with floating-point numbers, so the choice rarely matters.

Detaching Gradients

Sometimes you want to stop gradients from flowing:

#![allow(unused)]
fn main() {
let x = Tensor::from_vec(vec![2.0], true);
let y = &x * &x;  // y = x²

// Detach: treat y as a constant for further operations
let y_detached = Tensor::from_vec(y.data().clone(), false);  // requires_grad=false

let z = &y_detached + &x;  // z = y_detached + x (y treated as constant)

backward(&z);

println!("dz/dx = {}", x.grad()[0]);  // 1.0 (only from addition, not from y)
}

Use case: Stopping gradient flow in certain model parts (e.g., frozen layers).

In-Place Operations Warning

In-place modifications break the computational graph:

#![allow(unused)]
fn main() {
let mut x = Tensor::from_vec(vec![1.0, 2.0], true);
let y = &x * &x;

// ❌ BAD: Modify x in-place
x.data_mut()[0] = 5.0;

backward(&y);  // ⚠️ Undefined behavior! x changed after being used
}

Solution: Entrenar prevents in-place modifications for tensors with requires_grad=true:

#![allow(unused)]
fn main() {
// Entrenar's safeguard
if x.requires_grad() {
    panic!("Cannot modify tensor with requires_grad=true in-place");
}
}

Computational Complexity

Operation	Forward	Backward	Total
Add/Mul	O(n)	O(n)	O(n)
MatMul	O(mnk)	O(mnk)	O(mnk)
ReLU	O(n)	O(n)	O(n)
LayerNorm	O(n)	O(n)	O(n)
Attention	O(n²d)	O(n²d)	O(n²d)

Key insight: Backward pass has same asymptotic complexity as forward pass.

Debugging Gradients

Check if Gradients are Computed

#![allow(unused)]
fn main() {
let x = Tensor::from_vec(vec![2.0], true);
let y = &x * &x;

backward(&y);

if x.grad()[0] == 0.0 {
    eprintln!("Warning: Gradient is zero (might indicate issue)");
}
}

Gradient Explosion/Vanishing

#![allow(unused)]
fn main() {
fn check_gradients(params: &[&Tensor]) {
    for param in params {
        let grad_norm = param.grad().iter().map(|g| g * g).sum::<f32>().sqrt();

        if grad_norm > 100.0 {
            eprintln!("Warning: Gradient explosion (norm={})", grad_norm);
        } else if grad_norm < 1e-7 {
            eprintln!("Warning: Gradient vanishing (norm={})", grad_norm);
        }
    }
}
}

Gradient Checking

Always validate custom operations with finite differences:

#![allow(unused)]
fn main() {
#[test]
fn test_my_operation_backward() {
    let x = Tensor::from_vec(vec![1.0, 2.0, 3.0], true);
    let y = my_custom_operation(&x);

    backward(&y);

    // Compare with numerical gradient
    check_gradient(&y, &x, epsilon=1e-3, threshold=0.2);
}
}

Key Takeaways

Backward pass applies chain rule in reverse topological order
Each operation implements local gradient rule (e.g., mul: dx = y*dz)
Gradients accumulate when tensors have multiple consumers
Matrix operations use transposition for gradient computation
Nonlinear activations use derivative of activation function
Normalization requires saved statistics from forward pass
Complexity of backward equals forward (asymptotically)

What's Next?

Gradient Computation - Mathematical derivations
Finite Difference Validation - Testing gradients
Tensor Operations - All supported operations

Ready to dive into the math? Continue to Gradient Computation →

Entrenar - Training & Optimization Library