A C++ Graph-based autodiff engine for neural networks

The four models in this series so far share a property that is easy to miss: none of them has a loss function. The Hopfield network stores patterns with Hebb's rule. The restricted Boltzmann machine subtracts two correlations. The Kohonen map moves the winner and its neighbours. The reservoir readout fits a linear map. In every case the update rule was derived once, on paper, and then written into the code by hand. Change what you want the model to optimise and you rewrite the arithmetic.

That is the boundary this post crosses. The repository contains a small automatic differentiation engine in src/math/autograd, and it is the component that lets the user write the objective and lets the machine produce the gradient. It is also what the multilayer perceptron in the reservoir post trains against.

The engine takes a specific approach, and the approach is the interesting part. It does not record operations while the program runs and replay them backwards. It builds the objective as an explicit graph of nodes, and then differentiates that graph into a second graph. Both graphs are then compiled into flat evaluation orders and run millions of times. The work of the chain rule happens once, at setup, rather than once per training sample.

An expression is a graph of typed nodes

Everything rests on one small class. An ExpressionNode carries a type tag, its children, and its shape:

class ExpressionNode {
public:
    const ExpressionType type;
    // Note the reference to its own type, allowed because its just a pointer
    std::vector<ExpressionNode*> children;
    dim_t row;
    dim_t col;

    // Anonymous union to contain all the node data, this is a
    // bit corny but reduces storage (which is significant when handling
    // images)
    union {
        float vector_constant;
        unsigned long expr_id;
        unsigned long lp_norm;
        unsigned long pow;
    } expression_data;

The union is worth pausing on. A node never needs more than one scalar payload: a constant needs its value, a power needs its exponent, a norm needs its p. Storing them in a union keeps every node the same small size regardless of type. The comment about images is the real motivation. A constant vector of 784 ones stores a single float and its column count, not 784 floats, because ExpressionType::Constant means "this value, repeated across the whole vector".

Nodes are never constructed directly. A NodeGenerator owns them all in one arena:

class NodeGenerator {
    std::vector<std::unique_ptr<ExpressionNode>> nodes;
public:
    // Explicitly disallow copies, we are stories std::unique_ptr so
    // compiler forbids us from having copy operators!
    NodeGenerator(const NodeGenerator&) = delete;

This gives the graph a single owner with a single lifetime. Every node that belongs to one objective dies together when the generator does, and no node in the graph ever has to know who else points at it. That matters more than it first appears, because the graphs in this engine are not trees.

The generator is also where shape checking lives. Operations refuse mismatched operands when the node is built, not when it is evaluated:

inline ExpressionNode* sum(ExpressionNode* a, ExpressionNode* b) {
    if (a->cols() != b->cols() || a->rows() != b->rows())
        throw std::runtime_error("Cannot sum elements of different size");
    return create_vector_op(ExpressionType::Add, { a, b }, a->cols(), a->rows());
}

Because the graph is built once and evaluated many times, a shape error surfaces during setup instead of in the middle of the training loop.

Writing a loss and getting a graph

The MLP example builds mean squared error. Two variables, a subtraction, a squared norm:

ScalarFunction loss_function;
auto& g = loss_function.generator();
auto y = g.create_vector_variable(100);
auto y_hat = g.create_vector_variable(100);
loss_function = g.squared_l2_norm(g.sub(y, y_hat));

A note on the names before we go further: in this example y is the variable that the network's output is bound to, and y_hat receives the ground truth. The trainer differentiates with respect to y. I keep the identifiers from the source below.

squared_l2_norm is not a node type. It is two nodes, composed by the generator:

ExpressionNode* squared_l2_norm(ExpressionNode* expr) {
    auto* lp = lp_norm(expr, 2);
    return pow(lp, 2);
}

So the loss becomes a chain of four nodes over two variables. The shapes tell the story: the two variables and the subtraction carry 100 columns, and the norm collapses them to a single column that the power then keeps.

The right-hand graph is the output of the engine, and the arrows converging on Sub are the point of this post. The derivative graph does not copy the forward graph. It reuses its nodes.

Differentiation builds a second graph

ScalarFunction::derivative takes an empty VectorFunction and fills it in:

void derivative(VectorFunction<ScalarType>& func, VectorVar var, unsigned int optimize_steps = 10) {
    // Compute the derivative of this scalar function with respect to a variable
    Differentiator diff;
    VectorExpression expr = diff.differentiate<ScalarType>(root, var, func.generator());
    func = expr;
}

The types change on purpose. A ScalarFunction has a root of one column; its derivative with respect to a vector variable has as many columns as that variable, so it is a VectorFunction. The generator passed to differentiate belongs to the derivative, which means the new nodes are owned by the gradient while the nodes they reuse stay owned by the loss.

The traversal runs on an explicit stack with a cache keyed on the node pointer:

struct DiffFrame {
    ExpressionNode* node;
    std::vector<ExpressionNode*> dchildren;
    bool ready;
};

if (!frame.ready) {
    frame.ready = true;
    // Push children
    for (auto* child : node->get_children()) {
        if (!cache.count(child)) {
            st.push({ child, {}, false });
        }
    }
}
else {
    // Fetch all derivatives from the mapping.
    frame.dchildren.clear();
    for (auto* child : node->get_children())
        frame.dchildren.push_back(cache[child]);
    ExpressionNode* dnode = differentiate_node(node, frame.dchildren, var, gen);
    cache[node] = dnode;
    st.pop();
}

Each frame is visited twice. The first visit pushes the children; the second visit reads their finished derivatives out of the cache and builds the node's own. The cache does two jobs at once. It guarantees that every node is differentiated once even when several parents share it, and it makes the traversal safe on a graph rather than only on a tree. Avoiding recursion is deliberate as well: a deep objective would otherwise put its depth on the C++ call stack.

The rules themselves are a flat dispatch on the node type, and each one emits nodes instead of computing numbers. The product rule is the clearest example:

else if (node->get_type() == ExpressionType::ScalarMultiply) {
    auto left_d = dchildren[0];
    auto left = node->get_children()[0];
    auto right_d = dchildren[1];
    auto right = node->get_children()[1];

    return gen.sum(
        gen.multiply(left_d, right),
        gen.multiply(left, right_d)
    );
}

Read gen.multiply(left_d, right) carefully. right is a node from the forward graph, inserted directly into the derivative graph as a child. This is where the sharing in the figure comes from. The chain rule needs the value of a subexpression, and rather than rebuilding that subexpression, the engine points at the one that already exists.

The power rule composes the same way, building three nodes out of one:

else if (node->get_type() == ExpressionType::ScalarPower) {
    if (node->get_pow() == 1)
        return dchildren[0];
    else
        return gen.multiply(
            gen.multiply(
                gen.pow(node->get_children()[0], node->get_pow() - 1),
                node->get_pow()
            ), dchildren[0]);
}

For the squared norm, get_pow() - 1 is 1, and gen.pow returns its argument unchanged in that case, so the exponent node never appears. What remains is the LpNorm multiplied by the constant 2 and then by the derivative of the norm, which is exactly the ScalarMul and VecScalMul pair in the figure.

Simplification pays off because losses are full of constants

Derivative rules generate zeros. Every constant contributes one, and an objective written by a human is full of constants: targets, coefficients, thresholds. Left alone, those zeros would be multiplied and added at every evaluation for the rest of training. So differentiate ends with a rewrite pass that runs until it stops changing anything:

unsigned int performed_opts;
do {
    performed_opts = 0;
    optimize(der, gen, performed_opts);
} while (performed_opts > 0);

The rewrites are local algebraic identities, applied bottom-up on another explicit stack:

if (node->get_type() == ExpressionType::VectorScalarMultiply) {
    if (children[0]->is_zero() || children[1]->is_zero())
        return gen.zero(children[0]->cols(), children[1]->rows());
}
if (node->get_type() == ExpressionType::Add) {
    if (children[0]->is_zero() && children[1]->is_zero())
        return gen.zero(children[0]->cols(), children[0]->rows());
    if (children[0]->is_zero() && !children[1]->is_zero())
        return children[1];
    if (!children[0]->is_zero() && children[1]->is_zero())
        return children[0];
}

The second demo in main.cpp shows how much this removes. It builds a ridge-regularised objective over a 30-dimensional variable:

ScalarFunction func; // a scalar function
auto& g = func.generator();
auto u = g.create_vector_variable(30);
const double lambda = 0.01;
const auto expr = g.sum(
    g.squared_l2_norm((g.sub(u, 1.0))),
    g.multiply(lambda, g.squared_l2_norm(u))
);
func = expr;

Differentiating this with respect to u produces three zeros before the rewrite pass runs. The constant vector of ones inside the subtraction gives one. The constant λ gives another, which then multiplies a whole branch of the product rule. The third comes from that product itself. The rewrites collapse all three, and the surviving structure has one branch per term of the objective:

Sum[30]
|-- VectorScalarMul[30]              d/du of ||u - 1||^2
|   |-- ScalarMul[1]
|   |   |-- LpNorm[1]  p=2   ------> Sub[30] (shared with the loss graph)
|   |   ^-- Const[1]   v=2
|   ^-- LpNormDer[30]  p=2   ------> Sub[30] (the same node again)
^-- VectorScalarMul[30]              d/du of lambda*||u||^2
    |-- Const[1] v=0.01
    ^-- VectorScalarMul[30]
        |-- ScalarMul[1]
        |   |-- LpNorm[1] p=2 -----> Var u[30]
        |   ^-- Const[1]  v=2
        ^-- LpNormDer[30] p=2 -----> Var u[30]

This is the shape of the graph, not a transcript of the program's output; print_as_binary_tree prints the same structure with node addresses and dimensions. Two things stand out. The product rule for λ × ‖u‖² has become a single scaled branch, because the term in which λ was differentiated folded away. And Var u appears three times in the picture but exists once in memory, shared between the loss graph and both branches of its gradient.

Compiling the graph into a flat order

A graph is a poor thing to walk in an inner loop. Assignment therefore does more than store the root:

ScalarFunction& operator= (ScalarExpression expr) {
    if (expr != nullptr && expr->is_2d() || expr->cols() != 1)
        throw std::runtime_error("Cannot generate a scalar function from a vector expression!");
    root = expr;
    // This compiles the expression into a well defined evaluation ordering.
    evaluate_graph_non_recursive_ordering(expr, this->eval_ordering);
    return *this;
}

The ordering routine is a post-order traversal with a visited set, producing a vector in which every node appears after its children and exactly once. Evaluation then becomes a straight loop over that vector with a switch on the node type, writing each result into a cache:

for (ExpressionNode* node : eval_ordering) {
    switch (node->get_type()) {
    case ExpressionType::Sub: {
        auto* a = node->get_children()[0];
        auto* b = node->get_children()[1];
        cache[node] = cache[a] - cache[b];
        break;
    }

The cache is the reason the sharing in the derivative graph turns into real savings. Sub appears under both LpNorm and LpNormDer, but it appears once in the ordering, so y − ŷ is computed once per evaluation and read twice. The comment in apply states the consequence plainly:

// Note: WE CANNOT ASSUME no-aliasing of the data vectors because the expression tree is
// really an expression DAG because of memoization!

DAG here means directed acyclic graph: a node may have several parents, but no path leads back to itself. That is precisely what the derivative rules build, and it is why the flat ordering has to be computed with a visited set rather than by walking children twice.

Binding variables without copying

Evaluation needs the numbers for the variables. They arrive through a map of references, not of values:

// Raw buffer from eigen vector.
template <typename ScalarType>
using EvalMap = std::map<ExpressionNode*,
    std::reference_wrapper<EigVec<ScalarType>>>;

and the result is written through an Eigen::Ref, which binds to memory the caller already owns:

void apply(const EvalMap<ScalarType>& eval_data, Eigen::Ref<EigVec<ScalarType>> vec) {

Together these two choices let a gradient be evaluated straight into a column of an existing matrix. That is exactly how the trainer uses it, and it is the last piece needed before the training loop makes sense.

Where the engine sits in a training step

The trainer differentiates the loss once, when it is given one:

void set_loss_function(LossFunction* loss, LastOutput var, LastOutput ground_truth) {
    loss_function = loss;
    loss_function->derivative(gradient_generator, var);
    layer_variable = var;
    ground_truth_variable = ground_truth;
}

After that call, gradient_generator holds a compiled graph, and no differentiation happens again for the rest of training. The per-batch work is binding and evaluating:

// Then we construct the batched gradient matrix
timer.start("gradient");
for (size_t i = 0; i < batch.size(); ++i) {
    VectorType out_col = batch_output.col(i);
    VectorType true_col = batch.y_of(i);
    loss_eval_map.clear();

    loss_eval_map.emplace(layer_variable, out_col);
    loss_eval_map.emplace(ground_truth_variable, true_col);
    gradient_generator(loss_eval_map, batch_gradient.col(i));
}
timer.stop("gradient");

// We then perform backpropagation
timer.start("backprop");
network.backward(batch_gradient);

Notice what the engine is asked for, and what it is not asked for. It produces batch_gradient, the derivative of the loss with respect to the network's output. Nothing more. The derivatives with respect to the weights come from MultiLayerPerceptron::backward, which is ordinary matrix backpropagation and takes that gradient as its seed:

void backward(const Matrix& dLoss_dA_L) {
    int batch_size = dLoss_dA_L.cols();

    // Output layer delta
    delta[num_layers - 1] =
        dLoss_dA_L.cwiseProduct(act_functions[num_layers - 2].df(Z[num_layers - 1]));

    // Hidden layers
    for (int l = num_layers - 2; l >= 1; --l) {
        delta[l].noalias() =
            (W[l + 1].transpose() * delta[l + 1])
            .cwiseProduct(act_functions[l - 1].df(Z[l]));
    }

Why this division of labour suits neural networks

A trained network has two very different kinds of derivative in it, and this design gives each one to the tool that handles it best.

The layer gradients are regular. Every layer computes the same operation, Wa + b followed by an elementwise nonlinearity, and its derivative is always a transposed matrix product and an elementwise multiply. The structure never varies, the matrices are large, and the work belongs in W[l + 1].transpose() * delta[l + 1], where Eigen dispatches to a blocked kernel. There is nothing for a symbolic engine to discover here, and a general graph would stand between the training loop and the BLAS call.

The loss is the opposite. It is small, it is one vector wide, and it is the part the user actually changes: mean squared error today, cross-entropy tomorrow, either with a penalty term added. Writing its gradient by hand is where sign errors and missing factors come from, and it is the part that has to be rederived every time the objective changes. That is the work worth automating, and it is what the engine automates.

Three properties of the implementation follow from this split and are worth naming.

The setup cost is paid once. Building the graph, differentiating it, simplifying it and computing two evaluation orders all happen before the first batch. A tape-based engine records the operations of every forward pass and walks the recording backwards, which means a fresh graph for every sample. Here the graph is a fixed object and the loop only pushes numbers through it. For a network that sees hundreds of thousands of samples, that is the right place to put the work.

The gradient graph reuses the loss graph. Backpropagation needs forward values to compute backward ones, and this engine gets that by construction rather than by bookkeeping: the derivative rules insert the forward nodes themselves as children. The evaluation cache then computes each shared node once. In the MSE graph, y − ŷ is computed once and used by both the norm and its derivative.

Fused nodes carry the good gradient. The engine has a node for softmax cross-entropy taken as a unit, and its derivative rule is one line:

else if (node->get_type() == ExpressionType::SoftMaxCrossEntropy) {
    auto logits = node->get_children()[0];
    auto ref_val = node->get_children()[1];
    return gen.sub(gen.softmax(logits), ref_val);
}

The derivative of the cross-entropy of a softmax with respect to the logits is the softmax minus the one-hot target. Composing that gradient out of an exponential, a sum and a logarithm would produce a much larger graph passing through the softmax Jacobian. Declaring the pair as one node gives the classifier's gradient as a single vector subtraction. The reservoir readout uses exactly this node:

ScalarFunction loss_function;
auto y = loss_function.generator().create_vector_variable(5);
auto y_hat = loss_function.generator().create_vector_variable(5);
loss_function = loss_function.generator().smce_logits_true(y, y_hat);

The forward evaluation of that node subtracts the maximum logit before exponentiating, the standard guard against overflow:

EigVec<ScalarType> softmaxed = (logits_val.array() - logits_val.maxCoeff()).exp().matrix();
softmaxed /= softmaxed.sum();
const ScalarType scalar_loss = -(true_val.array() * (softmaxed.array() + eps).log()).sum();

Adding a new loss to the engine is therefore a contained job: an entry in ExpressionType, a factory method on the generator, a case in differentiate_node and a case in eval_node. The training loop does not change, because it only ever asks for the derivative of whatever root it was handed.

The line between derived and written

This module is where the repository stops encoding what a model should learn and starts deriving it. Hebb's rule, contrastive divergence and the Kohonen update are answers that somebody worked out and then transcribed. A graph plus a differentiator is a machine for producing such answers, and the price of a new objective drops from a derivation to a few lines of construction.

The design also says something about where automatic differentiation earns its cost. It is not a replacement for knowing the mathematics of your layers; the MLP's backward pass is still hand-written, and it is faster and clearer for being so. It is a tool for the part of a network that a researcher rewrites most often and reasons about least reliably. Drawing that line deliberately, rather than putting every operation behind the same abstraction, is what makes an engine this small useful at all.

There is a pleasing symmetry in where the series ends up. The Hopfield network in the first post computes with a graph whose edges are weights, and it settles into a minimum of an energy it cannot change. Here the graph's edges are operations, the minimum belongs to a function the user wrote, and the descent direction is produced by differentiating the graph itself.