Colored Markov polycategories and diagrammatic differentiation

TL;DR

This paper introduces a categorical framework using colored Markov polycategories for modeling, composing, and differentiating typed stochastic systems with complex wiring.

Contribution

It develops a novel colored Markov polycategory structure with trace semantics and diagrammatic differentiation for finite acyclic stochastic diagrams.

Findings

Proves structural laws for kernel slotwise composition.

Defines colored Markov polycategories with typed connections.

Derives local reverse-mode gradients for parameterized stochastic systems.

Abstract

Many stochastic systems are built by wiring typed components together, but the wiring is often neither purely sequential nor type-homogeneous. This paper develops categorical semantics for such systems using ordered polycategories whose morphisms are Markov kernels. The basic operation is kernel slotwise composition, which connects one output slot of a many-output kernel to one input slot of another and marginalizes the internal wire. We prove its structural laws by assigning trace semantics to finite acyclic diagrams. We then introduce colored Markov polycategories, where objects and kernels carry colors and typed connections are realized by coherent interface kernels. This gives a colored kernel slotwise composition and trace semantics for typed stochastic diagrams. To describe systems whose structure changes, we co-index colored Markov polycategories and parameter spaces over an indexing category. Finally, for finite acyclic parameterized diagrams, we prove a diagrammatic differentiation result. The derivative of an expected scalar objective is obtained from local reverse-mode contributions at the parameterized vertices, with stochastic and deterministic kernels handled through admissible local gradient operators. The construction gives a typed, compositional language for finite acyclic stochastic systems and their parameter sensitivities.