Automatic Differentiation on a Microcontroller: How numx Does It

Automatic differentiation was supposed to require a garbage-collected runtime with dynamic computation graphs. Most implementations depend on heap allocation to build and traverse those graphs. numx proves otherwise.

What automatic differentiation actually is

Automatic differentiation is not numerical differentiation (finite differences) and it is not symbolic differentiation (a computer algebra system). It is a third thing: exact computation of derivatives by applying the chain rule mechanically to the sequence of arithmetic operations in your program.

Given a function f, automatic differentiation gives you the exact derivative to floating-point precision, not an approximation, not a symbolic expression, but the number. This is what makes it useful for optimization: gradient descent needs exact gradients, and finite differences are both inaccurate and expensive.

Forward mode: dual numbers

Forward mode works by extending the number system with a dual component. A dual number has the form a + b*epsilon, where epsilon^2 = 0. The real part tracks the function value; the dual part tracks the derivative.

Every arithmetic operation on dual numbers propagates both components according to the usual calculus rules. Addition: (a + b*eps) + (c + d*eps) = (a+c) + (b+d)*eps. Multiplication: (a + b*eps) * (c + d*eps) = a*c + (a*d + b*c)*eps.

To compute f'(x) with forward mode, you evaluate f with input x + 1*epsilon. The real part of the output is f(x); the dual part is f'(x).

In numx, this looks like:

#include "numx/numx.h"
static numx_dual_t my_func(numx_dual_t x) {
    /* f(x) = x^3 + 2*x */
    numx_dual_t x2 = numx_dual_mul(x, x);
    numx_dual_t x3 = numx_dual_mul(x2, x);
    numx_dual_t two = numx_dual_const(2.0f);
    numx_dual_t lin = numx_dual_mul(two, x);
    return numx_dual_add(x3, lin);
}
int main(void) {
    numx_real_t val, deriv;
    numx_status_t s = numx_ad_dual_eval(my_func, 3.0f, &val, &deriv);
    /* val == 33.0, deriv == 29.0 (3*x^2 + 2 at x=3) */
    return s != NUMX_OK ? 1 : 0;
}

No allocation. The dual number is two floats on the stack. The derivative computation is as fast as the function evaluation itself.

Forward mode is ideal when you have one input and many outputs, or when the function has a single scalar input.

Reverse mode: static tape

Reverse mode is the algorithm behind backpropagation in neural networks. It is efficient when you have many inputs and a single scalar output, which is exactly the gradient-of-a-loss-function use case.

The standard implementation builds a dynamic computation graph (a directed acyclic graph of operations) on the heap during the forward pass, then traverses it backwards. This is why most autodiff libraries require garbage collection.

numx uses a static tape instead. The tape is a fixed-size array of nodes, allocated by the caller at compile time. The forward pass records operations onto the tape. The backward pass traverses the tape in reverse.

#include "numx/numx.h"
#define TAPE_CAPACITY 64
static numx_tape_node_t tape_buf[TAPE_CAPACITY];
static numx_tape_t tape;
int main(void) {
    numx_ad_tape_init(&tape, tape_buf, TAPE_CAPACITY);
    /* Create leaf variables */
    numx_ad_var_t x = numx_ad_var(&tape, 2.0f);
    numx_ad_var_t y = numx_ad_var(&tape, 3.0f);
    /* f(x, y) = x^2 * y + y */
    numx_ad_var_t x2   = numx_ad_mul(&tape, x, x);
    numx_ad_var_t x2y  = numx_ad_mul(&tape, x2, y);
    numx_ad_var_t loss = numx_ad_add(&tape, x2y, y);
    /* Backward pass */
    numx_status_t s = numx_ad_backward(&tape, loss);
    if (s != NUMX_OK) return 1;
    /* Exact gradients: df/dx = 2*x*y = 12, df/dy = x^2 + 1 = 5 */
    numx_real_t dx = numx_ad_grad(&tape, x);
    numx_real_t dy = numx_ad_grad(&tape, y);
    return 0;
}

The tape_buf is a static array. You size it at compile time based on your computation graph's depth. The library tells you at runtime if you exceed capacity (NUMX_ERR_OVERFLOW). No heap, no garbage collector, exact gradients on a chip with 512 KB of RAM.

Why this matters on embedded hardware

Embedded systems increasingly run model inference and lightweight on-device learning. A sensor node that can update a simple model from local observations, without a cloud round-trip, is qualitatively different from one that cannot.

numx's autodiff module makes gradient-based optimization possible in environments where it was previously infeasible: a microcontroller running a Kalman filter that can refine its noise parameters from observations, a motor controller that can fine-tune its PID coefficients from error signals, an edge sensor that can adapt a simple classifier from labeled examples collected locally.

The constraint is the tape size. A TAPE_CAPACITY of 64 nodes handles a few layers of a small neural network. More capacity requires more RAM. On an ESP32 with 320 KB of internal SRAM, a tape of 1024 nodes is entirely feasible.

The math that was supposed to require a server now runs on hardware smaller than a coin.