← Back to Projects
Machine LearningIteratingFeatured

MiniTorch-OCaml

Reverse-mode autodiff in OCaml: forward builds a graph, backward fills grads, gradcheck keeps me honest.

Framework backprop always felt like a black box, so I wrote my own. Nodes store values, grads, parents, and ops. Forward grows the graph; backward walks it. Most bugs do not show up in the forward pass; they show up when a shared node silently gets the wrong accumulated gradient. Finite-difference gradcheck is what I trust. The tiny MLP with SGD/Adam is just a smoke test that the pieces still talk to each other.

Engineering highlights

  • Graph nodes with values, grads, parents, and op metadata for reverse-mode
  • Backward rules for arithmetic, activations, reductions, transpose, and matmul
  • Central-difference gradchecks before I believe any training curve
Date
Spring 2026
Focus
ML Systems
Build stage
Extending operator coverage
Disciplines
Graph Construction · Reverse-Mode AD · Numerical Validation
Source code
ml-systemsautodiffocamlbackpropgradcheck
Autodiff visual
Full write-up below. The hero is the short version. This is the build story: what I designed, what broke, and why I made the calls I did.

Motivation

Why I started this

What pulled me in, and what I wanted to get better at.

Why I built it

I was tired of treating PyTorch as a calculator I did not understand. If backprop is the whole game, I wanted to implement it.

What interested me

Calling backward() is easy. Keeping parent links and local derivatives correct when a node is reused is not. Wrong grads still train for a bit, which is worse.

What I wanted to learn

OCaml made the graph nodes explicit. Gradcheck against finite differences is what actually caught my calculus mistakes.

System Overview

How the system fits together

MiniTorch is reverse-mode autodiff in OCaml. Forward ops grow a graph. Backward walks parents and fills grads. A tiny MLP with SGD/Adam is the smoke test. Gradcheck is the real test.

01

Tensor / graph nodes

Hold values, gradients, parents, and operation tags that define the local backward rule.

02

Forward ops

Elementwise arithmetic, activations, reductions, transpose, and matmul that extend the graph.

03

Reverse engine

Topological reverse traversal that applies local gradients and accumulates into parents.

04

Gradcheck

Central-difference numerical checks that catch calculus bugs types cannot see.

05

Training demo

Toy MLP with SGD/Adam showing loss decrease when autodiff is correct.

Data flow

Inputs → forward ops (graph grow) → scalar loss → reverse traversal → parameter gradients → optimizer step.

Control flow

Forward is eager graph building. Backward is an explicit reverse walk. Gradcheck samples perturbations around a point and compares analytic vs numerical derivatives.

Computation graph
MiniTorch computation graph overview

Each op records parents so reverse-mode can propagate local derivatives.

Engineering Breakdown

Broken down by discipline

Each block covers the goal, the design, what broke, what changed, and what shipped.

01Graph Construction

Graph Construction

Goal

Make every forward operation an explicit node that records enough metadata for a correct local backward rule.

Design

Typed graph nodes store value, grad, parents, and op. Arithmetic, activations (ReLU, Tanh, Exp, Log, Sigmoid, Pow), Sum, Transpose, and MatMul each extend the graph during the forward pass instead of mutating values in place without history.

Challenges

  • API surface had to stay small while remaining extensible for new ops.
  • Shape and broadcasting mistakes look like training instability rather than type errors.
  • Forgetting a parent edge produces silent zero gradients.

Iterations

  • Scalar-friendly arithmetic nodes.
  • Activations with simple local derivatives.
  • Reductions and matmul once multi-dimensional training demos mattered.

Final implementation

A forward API that grows a directed graph suitable for reverse-mode accumulation without hiding parent structure.

Forward pass
Forward pass building the computation graph

Ops append nodes; values alone are not enough for backward.

02Reverse-Mode AD

Reverse-Mode AD

Goal

Propagate gradients from a scalar loss back to parameters with correct local Jacobians.

Design

Seed the loss gradient, walk parents in reverse, and apply each op's local rule. Accumulation must respect multiple consumers of the same node. This is the systems core: autodiff as graph algorithm plus calculus tables.

Challenges

  • Gradient bugs rarely show up in forward outputs.
  • Shared subgraphs require careful accumulation, not overwrite.
  • Matmul and reductions have local rules that are easy to transpose wrong.

Iterations

  • Elementwise reverse rules.
  • Activation derivatives.
  • Reduction and matmul reverse paths used by the MLP demo.

Final implementation

A reverse engine that fills .grad fields from parent references so optimizers see usable parameter updates.

Backprop
Reverse-mode gradient propagation

Reverse traversal applies local rules and accumulates into parents.

03Numerical Validation

Numerical Validation

Goal

Prove analytic gradients against finite differences before trusting a training curve.

Design

Central-difference gradcheck in the demo harness compares engine gradients to numerical estimates. New operators are expected to ship with forward, backward, and a check. Types catch structure; numerics catch calculus.

Challenges

  • Finite differences are noisy near non-smooth points like ReLU kinks.
  • Passing a loss curve without gradcheck can still hide a wrong op that rarely activates.
  • The current harness prints check output; hardening toward asserted CI tests is future work.

Iterations

  • Manual derivative spot checks.
  • Central-difference sweeps for core ops.
  • Training-loop smoke tests as an integration signal, not a substitute for gradcheck.

Final implementation

A gradcheck workflow used whenever operators change, plus demo training that shows loss decrease when the engine is healthy.

Gradcheck
Gradient checking against finite differences

Numerical agreement is the proof types cannot provide.

Key Design Decisions

Calls I actually made

What else was on the table, what I picked, and why it still made sense once the hardware was real.

01

OCaml for an autodiff engine

The problem

Which language keeps graph nodes explicit without drowning in boilerplate?

Alternatives considered

  • Python/NumPy
  • C++
  • OCaml

Tradeoffs

Python is the ML default and soft on invariants. OCaml makes node variants and exhaustive matches hard to leave incomplete.

Why I chose this

OCaml. Exhaustiveness on ops is part of the safety story.

02

Reverse-mode versus forward-mode

The problem

Which AD mode matches neural training?

Alternatives considered

  • Forward-mode
  • Reverse-mode
  • Mixed

Tradeoffs

Forward-mode scales with inputs; reverse-mode scales with outputs. Training needs reverse.

Why I chose this

Reverse-mode. One scalar loss to many parameters is the target workload.

03

Eager graph building

The problem

When is the graph constructed?

Alternatives considered

  • Define-and-run static graph
  • Eager ops that record parents

Tradeoffs

Static graphs optimize earlier; eager graphs match how students debug PyTorch-style code.

Why I chose this

Eager recording. Debuggability beat premature graph compilers.

04

Gradcheck as the operator gate

The problem

How do you know a new op is correct?

Alternatives considered

  • Trust training loss
  • Unit tests on hand values only
  • Finite-difference gradcheck

Tradeoffs

Loss curves are necessary but insufficient. Gradcheck is slower and catches silent calculus bugs.

Why I chose this

Require forward + backward + gradcheck when adding operators.

Evolution

How it got here

Bench bring-up, CAD fits, soldering, and the demos in between. Not just the final photo.

  1. Stage 1

    Core tensor nodes

    Values, grads, parents, and arithmetic ops that prove the graph idea.

    No media for this milestone yet.

  2. Stage 2

    Activations and reductions

    Local derivatives beyond +, *, and the first non-elementwise paths.

    No media for this milestone yet.

  3. Stage 3

    Matmul + MLP demo

    Enough linear algebra for a tiny network and optimizer loop.

    No media for this milestone yet.

  4. Next

    Hardened tests and broader ops

    Asserting gradchecks in CI, richer broadcasting, and a vectorized backend once the operator set is stable.

    No media for this milestone yet.

Results & Validation

What held up

What worked in the end, what I can show for it, and where it's still limited.

Reverse-mode engine

lib/tensor.ml operator set and reverse traversal

Forward builds a graph; backward fills gradients for supported ops including matmul and common activations.

Numerical checks

Gradcheck path in the demo harness

Central-difference gradcheck exercises analytic gradients for the implemented ops.

Training smoke signal

Demo training loop in main.ml

Toy MLP with SGD/Adam shows loss decrease when autodiff and updates agree.

Photos and clips

The runs and stills that match the results above.

Validation
Gradcheck evidence visual

Finite differences keep reverse-mode honest.

Limitations

  • No GPU backend, convolutions, or production framework surface.
  • test/ is still thin; gradcheck prints today rather than a full asserted suite.
  • Broadcasting and operator coverage are intentionally incomplete.

Reflection

Looking back

What surprised me, what I'd redo, and questions I'm still chewing on.

What surprised me

  • Most bugs were local Jacobian mistakes that still produced plausible forward values.
  • Types prevented structural holes more than calculus holes.

What I would redesign

  • Asserted gradcheck CI for every op.
  • Clearer shape error messages at op construction time.
  • Separate pure math tables from graph mutation code.

Future improvements

  • Broader operator set with broadcasting rules.
  • Vectorized backend once correctness is boring.
  • Richer optimizers and serialization for experiment replay.

Questions that emerged

  • Where should broadcasting live: op nodes or a lowering pass?
  • How do frameworks keep reverse-mode fast without losing debuggability?