Expression Graph Autodiff in Go

Computation Graphs & ML Infrastructure advanced How AI Works ~30 min read

Overview

Expression-graph automatic differentiation turns mathematical code into a directed graph, attaches values and derivatives to graph nodes, and executes the graph with a virtual machine. Instead of writing gradient formulas by hand, the engineer builds the forward computation once and lets the graph runtime derive or propagate gradients for the variables of interest.

Gorgonia is a useful reference design because it exposes this mechanism directly in Go. A program creates scalar or tensor nodes inside an expression graph, composes operations such as addition and multiplication, binds runtime values with Let, asks for gradients with Grad or a differentiation-capable VM, and executes the resulting graph with a TapeMachine or LispMachine.

The transferable lesson is not the specific library API. The important pattern is the separation between graph construction, gradient graph construction, and graph execution. That separation is the same idea behind larger ML frameworks: a frontend records operations, autodiff augments the graph with derivative work, and a runtime executes a compiled or interpreted program over concrete tensor values.

Core Concepts

Expression Graphs as Mathematical ASTs

An expression graph is an abstract syntax tree for tensor math. Leaf nodes represent inputs, constants, and parameters. Interior nodes represent operations such as Add, Mul, matrix multiplication, activation functions, or reductions. The graph records data dependencies before any concrete numeric execution happens.

Values, Gradients, and With-Respect-To Nodes

Each variable node can hold a runtime value and, after differentiation, a gradient. A gradient request names a scalar cost node and a set of with-respect-to nodes. The autodiff system then determines how changes in each requested node affect the cost.

Automatic vs. Symbolic Differentiation

Automatic differentiation propagates derivatives through the executed graph by applying local chain-rule rules. Symbolic differentiation constructs explicit gradient nodes in the graph. Gorgonia exposes both approaches: a differentiation-capable machine can populate node gradients during execution, while Grad(cost, wrt...) adds gradient nodes that can be executed and inspected like ordinary graph outputs.

Key Concept: Autodiff Is Graph Transformation Plus Execution

The runtime does not need a hand-written derivative for the full model. It only needs derivative rules for primitive operations and a graph that captures the dependency structure. Complex gradients emerge by composing those primitive rules along graph edges.

How It Works

Build the Forward Graph

The first step is to represent the forward computation as graph nodes. This example models f(x, y, z) = (x + y) * z. The operation calls return new graph nodes rather than immediately calculating numeric values.

package main

import (
    "fmt"
    "log"

    "gorgonia.org/gorgonia"
)

func main() {
    g := gorgonia.NewGraph()

    x := gorgonia.NewScalar(g, gorgonia.Float64, gorgonia.WithName("x"))
    y := gorgonia.NewScalar(g, gorgonia.Float64, gorgonia.WithName("y"))
    z := gorgonia.NewScalar(g, gorgonia.Float64, gorgonia.WithName("z"))

    q, err := gorgonia.Add(x, y)
    if err != nil {
        log.Fatal(err)
    }
    result, err := gorgonia.Mul(q, z)
    if err != nil {
        log.Fatal(err)
    }

    gorgonia.Let(x, -2.0)
    gorgonia.Let(y, 5.0)
    gorgonia.Let(z, -4.0)

    grads, err := gorgonia.Grad(result, z, x, y)
    if err != nil {
        log.Fatal(err)
    }

    machine := gorgonia.NewTapeMachine(g)
    defer machine.Close()

    if err := machine.RunAll(); err != nil {
        log.Fatal(err)
    }

    fmt.Println("result:", result.Value())
    fmt.Println("df/dz:", grads[0].Value())
    fmt.Println("df/dx:", grads[1].Value())
    fmt.Println("df/dy:", grads[2].Value())
}

Add Gradient Nodes

Grad takes the scalar cost node and the variables to differentiate with respect to. In the example, the returned gradient nodes correspond to df/dz, df/dx, and df/dy in the same order as the requested variables. Because those gradients are graph nodes, they can be executed by the same machine as the forward computation.

Execute with a Graph VM

A TapeMachine precompiles a mostly static graph into a linear instruction tape, then runs those instructions sequentially. That design trades away dynamic graph mutation during execution, but it improves repeated execution because compilation and optimization happen before the hot loop. A LispMachine is more dynamic and can perform automatic differentiation during execution, but it is less suited to compile-once-run-many workloads.

Practical Applications

Training Loops in Go Services

Expression-graph autodiff lets Go services train or adapt small models without sending every experiment through a Python stack. A service can keep data ingestion, feature extraction, model execution, and optimizer steps in one compiled binary while still relying on graph-level derivative propagation.

Custom Differentiable Systems

Some systems are easier to express as mathematical programs than as prebuilt neural-network layers: optimization objectives, differentiable simulators, ranking losses, or constrained regression problems. A graph API gives engineers a way to assemble those objectives from primitive operations and differentiate the result.

Compiler and Runtime Education

Small graph runtimes are valuable for learning how larger ML systems work. The same concepts recur in production compilers: node identity, shape and type metadata, graph walks, gradient construction, instruction scheduling, buffer allocation, and runtime execution.

Implementation Patterns

Separate Model Construction from Value Binding

Build the graph once, then bind values before each execution. This keeps the mathematical structure stable while allowing different input batches, parameter values, or training examples to flow through the same compiled program.

type ScalarModel struct {
    graph  *gorgonia.ExprGraph
    x      *gorgonia.Node
    y      *gorgonia.Node
    z      *gorgonia.Node
    cost   *gorgonia.Node
    grads  gorgonia.Nodes
}

func NewScalarModel() (*ScalarModel, error) {
    g := gorgonia.NewGraph()
    x := gorgonia.NewScalar(g, gorgonia.Float64, gorgonia.WithName("x"))
    y := gorgonia.NewScalar(g, gorgonia.Float64, gorgonia.WithName("y"))
    z := gorgonia.NewScalar(g, gorgonia.Float64, gorgonia.WithName("z"))

    sum, err := gorgonia.Add(x, y)
    if err != nil {
        return nil, err
    }
    cost, err := gorgonia.Mul(sum, z)
    if err != nil {
        return nil, err
    }
    grads, err := gorgonia.Grad(cost, x, y, z)
    if err != nil {
        return nil, err
    }

    return &ScalarModel{graph: g, x: x, y: y, z: z, cost: cost, grads: grads}, nil
}

Choose the VM for the Graph Shape

Use a tape-style executor when the graph structure is known ahead of time and will be run many times. Use a more dynamic executor when graph construction and execution need to interleave. The wrong choice usually shows up as either recompilation overhead or awkward restrictions on dynamic behavior.

Expose Gradients as First-Class Outputs

When gradients are represented as nodes, they can be logged, tested, reused, or differentiated again. This is useful for debugging training failures and for implementing higher-order methods. It also makes gradient checks easier because the expected derivative is an ordinary output of the graph program.

Performance Considerations

Common Pitfalls

Further Reading