Category-Theory Deep Learning in Racket

Deep learning is usually taught as a pile of engineering tricks: activation functions, backprop, dropout, weight decay. This chapter takes the opposite route. We build a small but complete deep-learning framework in Racket whose design is dictated by category theory — the mathematics of structure and composition. Every component you would normally bolt on ad hoc (gradient descent, dropout, uncertainty estimates, invariance, ensembling) turns out to be a named categorical construction with precise laws.

The framework is a working implementation of the survey by Jia, Peng, Yang & Chen (2025), “Category-Theoretical and Topos-Theoretical Frameworks in Machine Learning”, Axioms 14(3):204 (https://doi.org/10.3390/axioms14030204). By the end of the chapter you will have ~285 lines of Racket that train an XOR network, implement dropout as a lens, estimate predictive uncertainty by Monte Carlo, run k-means as a colimit computation, classify via a topos’ subobject classifier, glue ensemble predictions like sheaf sections, and distill knowledge through a natural transformation.

A category is a collection of objects with morphisms (arrows) between them, where morphisms compose associatively and every object has an identity morphism. That sounds abstract, but it is exactly the structure machine learning already has:

• Networks are built by composing layers.
• Loss functions compose with models.
• Ensembles combine predictors.
• Architectures are distinguished by which transformations they commute with.

Category theory’s payoff is that when you recognize your construction as a known pattern (a lens, a Markov category, a colimit, a sheaf), you get its laws, its composition rules, and its proofs for free — rather than discovering each one by trial and error. The survey organizes categorical ML into five perspectives, and this chapter implements all five:

In the Para construction, a neural-network layer is a parametric morphism : it consumes an input from and produces an output in , but it does so using parameters living in a separate space . So a layer is really a map

To train such a layer we need to know how a small change in its output propagates back to its parameters and its input. That bidirectional pair — a forward map plus a pullback — is a lens. The forward pass of a layer is the “get” direction; the backward pass is the “put” direction:

The crucial theorem is that lenses compose. Chaining layers composes their forward passes, and backpropagation is nothing but composing the pullbacks in the reverse order — a covariant functor on the gradient category. The SGD update is then simply an endomorphism on the model object.

A Markov category is the right setting for stochastic computation. Its morphisms are stochastic kernels: given an input, they return a probability distribution over outputs. From a programming point of view, a stochastic morphism is just a function that samples.

Two ML staples are stochastic morphisms in disguise. Dropout multiplies activations by a Bernoulli mask sampled at train time — a stochastic lens whose forward pass samples the mask and whose backward pass reuses the same mask (the closed-optic requirement). Bayesian layers treat weights as random variables, sampling at every forward pass; the spread of many forward passes measures epistemic uncertainty.

A layer is equivariant under a group when for every : the symmetry of the input is preserved by the layer. Invariant maps go further and forget the symmetry entirely. The sum is invariant to any permutation of its inputs — it is the colimit (coproduct) over the set diagram, and it is the heart of set-based architectures like DeepSets. In the same spirit, k-means clustering is a colimit computation: each cluster centroid is the colimit (average) of the points in its cluster, and point-to-centroid assignment is the universal morphism.

A topos is a category rich enough to do logic inside it. Its key ingredient is the subobject classifier : a special object with a “true” arrow such that every subset is picked out by a unique characteristic morphism with . In machine learning, the sigmoid output of a binary classifier is such a characteristic morphism: it maps inputs into , and the decision boundary is .

A sheaf packages local data that must agree on overlaps. If you think of each model in an ensemble as an expert covering its own local context, then combining their predictions requires the sheaf condition: on shared context, experts must agree. Consistent sections glue into a global prediction; inconsistent ones reveal disagreement. Finally, carries a Heyting algebra (intuitionistic logic); on it collapses to , , .

Given two functors and , a natural transformation is a family of morphisms that commutes with every arrow — a structure-preserving way to convert one kind of computation into another. Knowledge distillation is exactly this: a small “student” model learns to reproduce the hidden representations of a large “teacher”. The adapter between teacher and student representations is .

Before we look at any parsing or training code, here is the data it processes. The framework stores each example as a pair .

XOR, the canonical non-linearly-separable problem. Four points in the unit square, where the target is 1 exactly when the two inputs differ:

Dropout demo. One five-vector that we will push through the dropout lens, plus the unit gradient we backpropagate through it:

Bayesian demo. A three-dimensional input for the single-output Bayesian layer, which we will query 200 times:

Cluster data. Nine points in the plane arranged in three obvious clusters: three near (1,1), three near (5,5), and three near (1,5):

Sheaf sections. Three “expert” predictions on overlapping contexts, two of which agree within tolerance and one of which does not:

All randomness is seeded () so every run reproduces the exact outputs shown later in this chapter.

The framework is a single file, deep_learning_category_theory.rkt, and uses only the standard Racket distribution. It declares its language and the three libraries it needs:

racket/list supplies list tools (shuffle, take, in-naturals), racket/format supplies ~r for rounded number formatting, and racket/random supplies random-seed so the demos are reproducible.

The heart of the framework is a representation of a layer as a lens. Three structs encode the categorical objects: layer-params (the parameter space ), layer-grads (a tangent vector at , i.e. a gradient), and model (the product of a list of layer parameter spaces). A model is just a list of layers.

The small vector utilities are deliberately simple: lists stand in for vectors, and matrices are lists of rows. matvec multiplies a matrix by a vector; outer builds the outer product that appears in the weight gradient ; matT-vec multiplies by the transpose for the input gradient .

forward-para is the lens itself. It computes pre-activations and activations , then returns two things: the activations and a pullback closure. The pullback is the categorical put-back. Given the upstream gradient from the layer above, it computes the local gradient , packs the parameter gradients into a layer-grads (the outer product for dW, for db), and propagates on to the previous layer. Because the pullback is a closure, it remembers — that is what makes composing lenses possible without storing every intermediate activation globally.

Now we can compose layers. network-forward walks the model’s layer list, threading the current input and stacking pullbacks; the final output is the scalar prediction (the XOR network has one output unit). Backpropagation, model-backward, is literally the pullbacks composed in reverse order.

Note the one subtlety in model-backward: the loss gradient is a scalar, so it is wrapped in a one-element list before entering the first pullback; every pullback thereafter exchanges lists (activations and gradients are vectors). This is the “covariant functor on the gradient category” from the theory section: -vectors flow backward through exactly the same structure the activations flowed forward through.

Training is then routine plumbing over this categorical spine: Glorot initialisation, network construction from an architecture spec, an SGD loop over epochs, and a predict function.

train folds one pass over the whole dataset per epoch: each example advances the model (acc-m) and accumulates loss (acc-loss). Lazy Racket trivia: the fold is a for/fold, which is tail-recursive because each iteration’s next model state is produced by train-step before the loop continues.

First, dropout as a stochastic lens. make-dropout-lens returns a function that samples a Bernoulli mask, scales by (inverted dropout keeps the expected activation constant), and attaches a pullback that reuses the saved mask. Both passes share the mask — that shared state is the stochastic lens’s “closed optic” structure.

The Bayesian layer stores the mean weight matrix and a fixed spread . Each forward pass draws a fresh weight matrix (via a Box-Muller transform), then computes a sigmoid activation exactly like a deterministic layer would. The Monte Carlo wrapper runs forward passes, returns their average as the predictive mean, and their standard deviation as a measure of epistemic uncertainty: high -estimate means the model is unsure about this input.

Permutation-invariant pooling is one line: (apply + xs) — the colimit over the set diagram, invariant because addition is commutative. The k-means blocks read the same categorical way: update-centroids computes each cluster centroid as the average (colimit) of its points, and nearest-centroid plays the universal assignment morphism.

update-centroids is worth a close read: for/list over the cluster index collects the points labelled into cluster; then (apply map list cluster) transposes the list of points into a list of coordinate columns, each averaged. An empty cluster (possible when is unlucky) falls back to the zero vector rather than crashing. The k-means loop alternates assignment and centroid update until centroids stop moving — the fixed point is the categorical colimit structure in action.

subobject-classify is the characteristic morphism realized: it evaluates the trained model to get the probability and thresholds at to name the subobject. sheaf-glue checks the sheaf condition pairwise between all sections within tolerance and, if all agree, returns the average — the glued global section. Failure to agree returns #f, the categorical “no gluing”. The four Heyting operations are the min/max/complement definitions from the theory section.

The distillation adapter is a linear map with sigmoid squashing: takes a teacher hidden representation and produces the student representation. Because it is applied uniformly to any vector, it is a family of components — exactly a natural transformation .

The module+ main block seeds the random generator, then runs one demo per perspective. It feeds the XOR dataset to the freshly built network, pushes a vector through a dropout lens, queries the Bayesian layer 200 times, pools two permutations of the same set, clusters the nine 2-D points, classifies the trained XOR model, glues sheaf sections, evaluates Heyting formulas, and distills a teacher’s hidden layer into a student representation.

The distillation demo does something worth noticing: instead of re-running the whole network, it invokes forward-para directly on the first two layers of the trained model to harvest the 4-dimensional hidden representation — the lens API pays for itself when you want access to intermediate structure, not just the final prediction.

From the directory containing the file, run:

No packages beyond the standard Racket distribution are needed. Because of random-seed 42 the output is fully reproducible and looks exactly like this:

Run the program and read it as five little experiments, each confirming one categorical story.

I — the Para spine learned XOR, and the loss tells the story. The epoch losses collapse from 1.173366 to 0.002327 to 0.000651 — a steady exponential-ish decay with no divergence, which is what you expect from compositional backprop on a well-posed problem. The real check is the prediction table: (0,1) and (1,0) produce ~0.99 (class 1), while (0,0) and (1,1) produce ~0.01 (class 0). XOR is not linearly separable, so a network that nails it proves the lens composition is actually computing two nonlinear feature layers correctly, not just memorizing thresholds. The small residual (0.0064 vs. a perfect 0) is sigmoid saturation slack, nothing more.

II — dropout and Bayesian layers behave as stochastic morphisms. Of the five activations, elements 3 and 5 were dropped to 0; the survivors were rescaled by : 1.0 \to 1.429, 2.0 \to 2.857, 4.0 \to 5.714. The backward pass returns the same mask multiplied by the same scale — (1.429 1.429 0 1.429 0) — which is the closed-optic requirement made visible: no gradient flows through dropped units, and surviving gradients are rescaled so updates stay unbiased. The Bayesian MC query returns mean 0.3055, std 0.0717: the sigmoid-squashed predictions spread by a standard deviation of ~0.07 around 0.31, which quantifies how much the weight noise actually matters for this input. A confident model would report a much smaller std.

III — invariance and colimits are visible in two lines. Both permutations of the set {1,3,5,2} pool to 11.0: the colimit ignores order by construction. K-means finds the structure we planted: centroids (1.03, 0.97), (1, 4.97), (4.97, 5.03) — within ~0.05 of the true cluster means (1,1), (1,5), (5,5). The labels (0 0 0 2 2 2 1 1 1) are exact: each of the three planted clusters is recovered with zero misassignments. The centroid-as-colimit reading is not a metaphor here; it is literally how the code computes them (average of the cluster’s points).

IV — the topos logic behaves exactly like set-theoretic logic. The trained classifier’s χ readings mirror the XOR table: high probabilities (membership) on the two positive points, low on the two negative ones, and class thresholds at χ⁻¹(0.5) agree with the raw predictions. The sheaf gluing is the most instructive line: A+B (0.72 vs. 0.68, difference 0.04 < 0.1) glue into 0.7 — the average — while A+C (0.72 vs. 0.91, difference 0.19 ≥ 0.1) refuse to glue: inconsistent experts are flagged rather than silently averaged, which is the sheaf condition protecting you from ensembling disagreeing models. The Heyting line demonstrates the [0,1]-valued logic: 0.8 ∧ 0.3 = 0.3, ¬0.8 = 0.1999... (float arithmetic shows its 0.2 as 0.19999999999999996 — harmless), and implication 0.8 ⇒ 0.3 = max(0.2, 0.3) = 0.3.

V — distillation compresses the teacher’s representation. The teacher hidden vector for input (1.0, 0.0) is 4-dimensional: (0.8588 0.0626 0.2403 0.0987) — mostly-concentrated feature activity with one dominant unit. The student adapter squeezes it to 2 dimensions: (0.3703 0.5728). The point of the demo is not that these numbers are meaningful on their own; it is that is a uniform map — apply the same adapter to any teacher hidden state and you get a comparable student state, which is the naturality condition making distillation a bona fide natural transformation rather than a hand-fitted regression.

Category theory gave us a vocabulary, and the vocabulary gave us structure. The five sections of the file are five named categorical constructions: layers are Para morphisms and lenses, dropout and Bayesian layers are stochastic morphisms in a Markov category, pooling and k-means compute colimits, classification is characteristic map into with sheaf-conditioned ensembling, and distillation is a natural transformation. Each one compressed into a few lines of Racket that would otherwise be an ad hoc engineering special case — and each one carries the laws of its construction for free.

The complete framework is about 285 lines and uses only the standard Racket distribution. To explore further: replace sigmoid with ReLU (and its derivative), try deeper architectures by editing the make-network spec, or widen the dropout lens into an input-space perturbator — the categorical scaffolding does not change, only the morphisms you plug in.