Understanding Autograd: The Engine Behind Deep Learning

Automatic differentiation (autograd) is the cornerstone of modern deep learning. It enables neural networks to learn by automatically computing gradients of complex functions-no manual calculus required.

What is Autograd?

Automatic Differentiation is NOT:

• ❌ Symbolic differentiation (like Wolfram Alpha does)
• ❌ Numerical differentiation (finite differences)
• ✅ It's a algorithmic technique that applies the chain rule automatically

Key Concept: Autograd builds a computational graph as your code executes, tracking every operation. Then it traverses this graph backwards to compute gradients efficiently.

Computational Graphs Explained

Think of a computational graph as a blueprint of your computation:

Nodes: Variables and operations Edges: Data flow between operations

Simple Example: f(x, y) = (x + y) × x

Input:  x=2, y=3
 
Graph:
   x=2    y=3
    \    /
     (+)      → a=5
      |  \
      |   x=2
      |  /
     (×)       → f=10

Forward Pass: Compute f(x, y) = 10 Backward Pass: Compute ∂f/∂x and ∂f/∂y using chain rule

The Chain Rule: Foundation of Backpropagation

For composite functions, the chain rule states:

If z = f(y) and y = g(x), then:

dz/dx = (dz/dy) × (dy/dx)

Multi-variable version (used in neural networks):

∂L/∂w = Σᵢ (∂L/∂yᵢ) × (∂yᵢ/∂w)

Where:

• L = Loss function (what we want to minimize)
• w = Weight parameter
• yᵢ = Intermediate values in the computation graph

PyTorch vs JAX: Two Approaches to Autograd

PyTorch: Define-by-Run (Dynamic Graphs)

• Graph is built as code executes
• Different graph for each forward pass
• Easier to debug (feels like normal Python)
• Perfect for variable-length sequences, conditional logic

JAX: JIT Compilation (Static Graphs)

• Graph is traced once, then compiled
• Same graph for all inputs (must be same shape)
• Blazing fast execution (XLA compilation)
• Better for production deployment

Common Pitfalls & Debugging Strategy

1. Detached Tensors (No Gradient Flow)

Problem: Operations that break the computational graph

x = torch.tensor([2.0], requires_grad=True)
y = x.detach()  # ❌ Gradient won't flow through y
z = y ** 2
z.backward()    # x.grad will be None!

Fix: Avoid .detach(), .numpy(), or .item() in the middle of computation

2. In-Place Operations

Problem: Modifying tensors in-place can corrupt gradients

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
x += 1  # ❌ In-place modification
y.backward()  # RuntimeError!

Fix: Use x = x + 1 instead of x += 1

3. Shape Mismatches

Problem: Broadcasting can hide shape errors until backprop

# Forward pass works (broadcasts automatically)
x = torch.randn(32, 10)      # Batch of 32 samples
w = torch.randn(10, requires_grad=True)  # ❌ Should be (10, 1)
y = x * w  # Works but gives wrong shape (32, 10)
# Backward pass crashes or gives wrong gradients

Fix: Always verify tensor shapes with .shape before and after operations

4. NaN/Inf in Gradients

Most common causes:

• Division by zero: 1 / (x - x)
• Log of zero/negative: log(0) or log(-1)
• Exploding gradients: Very large learning rates
• Numerical overflow: exp(1000)

Gradient Checking: Verifying Your Autograd

Always verify autograd implementation with numerical gradients:

Finite Difference Approximation:

f'(x) ≈ [f(x + ε) - f(x - ε)] / (2ε)

Where ε is a small value (e.g., 1e-5)

If autograd gradient ≈ numerical gradient (within 1e-5), you're good!

When to Use Manual Gradients

Most of the time, use autograd. But manual gradients are needed for:

• Custom CUDA kernels (low-level GPU operations)
• Non-differentiable operations (argmax, sampling)
• Memory-constrained scenarios (gradient checkpointing)

Mental Model: Autograd as a Recording System

Think of autograd as a video recorder:

1. Press Record (requires_grad=True)
2. Perform operations (forward pass) - everything is recorded
3. Play backwards (.backward()) - gradients computed from recording
4. Get the gradients (.grad) - extract what was computed

Summary

Key Takeaways:

1. Autograd builds a computational graph during forward pass
2. Backward pass applies chain rule automatically
3. PyTorch = dynamic (flexible), JAX = static (fast)
4. Common bugs: detached tensors, in-place ops, shape mismatches, NaN/Infs
5. Use anomaly detection and gradient checking for debugging
6. The chain rule is the mathematical foundation of everything

Remember: Master autograd debugging = Master deep learning implementation!

Master automatic differentiation (autograd) and computational graphs in PyTorch and JAX. Learn practical debugging techniques for gradients, shape mismatches, and numerical instabilities (NaNs/Infs).