Stochastic Neural Network Symmetrisation in Markov Categories
Rob Cornish
TL;DR
The paper addresses preserving and upgrading symmetry in neural networks when outputs may be stochastic. It reframes equivariance and symmetrisation within Markov categories, defining a general, compositional framework that upgrades $H$-equivariant morphisms to $G$-equivariant ones via restriction functors and coset actions. A key contribution is a two-step end-to-end methodology that uses $L_\\varphi R_\\varphi$ and a precomposition by $\\Gamma$ to construct symmetrisation procedures, with results guaranteeing stability and, under suitable conditions, surjectivity. The framework unifies and extends prior deterministic methods (canonicalisation, frame averaging) and introduces stochastic symmetrisation for Markov kernels, enabling exact sampling and data-augmentation-like behaviour without expensive averaging. Empirically, the approach yields competitive or superior performance on synthetic tasks and provides a principled, scalable way to enforce symmetry in complex ML systems while preserving probabilistic structure and interpretability.
Abstract
We consider the problem of symmetrising a neural network along a group homomorphism: given a homomorphism $\varphi : H \to G$, we would like a procedure that converts $H$-equivariant neural networks to $G$-equivariant ones. We formulate this in terms of Markov categories, which allows us to consider neural networks whose outputs may be stochastic, but with measure-theoretic details abstracted away. We obtain a flexible and compositional framework for symmetrisation that relies on minimal assumptions about the structure of the group and the underlying neural network architecture. Our approach recovers existing canonicalisation and averaging techniques for symmetrising deterministic models, and extends to provide a novel methodology for symmetrising stochastic models also. Beyond this, our findings also demonstrate the utility of Markov categories for addressing complex problems in machine learning in a conceptually clear yet mathematically precise way.
