Composable Inference: Modular Probabilistic Reasoning

• Composable inference is a modular framework that constructs scalable probabilistic reasoning systems by combining reusable inference components such as kernels and effect handlers via principled rules.
• It bridges statistical modeling, program transformations, and categorical formalisms to enable flexible, compositional designs in probabilistic programming and machine learning.
• This paradigm supports multi-agent, distributed, and meta-learning scenarios by orchestrating specialized inference modules for rapid prototyping and formal convergence guarantees.

Composable inference is the paradigm by which inference algorithms, probabilistic models, or reasoning mechanisms are constructed from modular, reusable, and combinable components, rather than as monolithic procedures. Composability can manifest in inference over statistical models, programming-language abstractions, automated reasoning systems, machine learning frameworks, and hardware, reflecting a broad unifying principle: individual inference modules—be they kernels, handlers, programs, operators, or agents—can be assembled, manipulated, and orchestrated using principled composition rules to perform scalable, flexible, and robust probabilistic reasoning.

1. Fundamental Principles and Formalisms

Multiple formalizations underpin composable inference. In statistical modeling, the concept is instantiated by composite Bayesian inference, where separate Bayesian "agents"—each with a marginal likelihood (or model) for a partition of the data—are combined into a coherent composite posterior by log-linear pooling, with nonnegative convex weights ensuring external Bayesianness and conservative inference. Formally, for disjoint data blocks $\{D_i\}_{i=1}^K$ and agent models $p_i(D_i \mid \theta)$, the composite likelihood reads

$$L_c(\theta\mid D) = \prod_{i=1}^K p_i(D_i \mid \theta)^{\lambda_i},\quad \lambda_i \ge 0,\, \sum_i \lambda_i = 1,$$

yielding the composite posterior $$\pi_c(\theta\mid D)\propto \pi_0(\theta) L_c(\theta\mid D)$$ (Roche, 2015).

This compositional view extends to the categorical and programming-language level, where inference is represented as program transformations, effect handlers, or functorial mappings in Markov (or monoidal) categories. In these settings, the composition of inference operations (e.g., conditional updates, MCMC kernels, expectation evaluations) is encoded as the composition of morphisms or handlers, subject to algebraic and operational laws that guarantee correctness, composability, and sometimes convergence (Zinkov et al., 2016, Moore et al., 2018, Braithwaite et al., 2023).

2. Compositionality in Probabilistic Computing and Program Transformations

In probabilistic programming and program transformation frameworks, composable inference is realized by constructing inference algorithms as pipelines of program-to-program transformations. Models and algorithms are both represented as symbolic programs or abstract syntax trees (ASTs), and each transformation (e.g., disintegration, conditioning, normalization, Metropolis–Hastings, Gibbs) preserves or modifies the meaning of the input program according to defined rules.

For example, in (Zinkov et al., 2016), core transformations include:

• Disintegration: Constructs conditional programs from joint programs, yielding kernels for posterior inference.
• Normalization: Converts unnormalized measures into probability distributions.
• MCMC kernel composition: Produces a sampler by composing a proposal with a target density, with algebraic simplification and symbolic expectation handling.
• Composition algebra: Arbitrary chains of these transformations yield complex, often hybrid, inference strategies.

Because all transformations emit valid programs in the core language, they may be chained, nested, or swapped, facilitating easy assembly of exact, variational, or MCMC-based inference schemes.

Similarly, Edward2 leverages algebraic effect handlers—handlers are nested, and their composition is defined via substitutional and operational rules. This yields stacked inference transformations, enabling hybrid inference, sequential algorithms, and new inference strategies via mere addition to the handler stack (Moore et al., 2018).

3. Modular Inference in Bayesian and Categorical Structures

Categorical and lens-based approaches formalize Bayesian inversion as a compositional operation on stochastic kernels or Markov categories. The "Bayesian lens" pattern expresses pairs $(f, f^\dag)$, where $f$ is a stochastic map and $f^\dag$ is its inversion for a given prior, ensuring the Bayes coherence law holds for composition:

$$(f, f^\dag), (g, g^\dag) \implies (g\circ f, f^\dag \circ g^\dag),$$

with explicit state-dependent composition reflecting the chain rule for Bayesian updating (Braithwaite et al., 2023, Smithe, 2021). Such compositional laws guarantee that inversion of a composed process is the composition of the component inversions, up to support considerations. These theoretical guarantees directly inform modular, type-driven, and scalable approaches to probabilistic inference.

In graphical representations, such as string diagrams for Bayesian networks and causality, compositional laws govern the normalization and conditioning operations. Boxes (normalization) and wires/bends (disintegration) satisfy algebraic rules permitting the stepwise extraction and elimination of normalization steps, culminating in open diagrams encoding the conditional or interventional distributions—thus enabling compositional, graphical inference on structured probabilistic models (Jacobs et al., 28 Nov 2025).

4. Composable Inference as Algorithmic and Systemic Design

Contemporary computational inference libraries and frameworks embody composable inference in their architectural design. BlackJAX (Cabezas et al., 2024) realizes Bayesian inference by exposing algorithmic "atoms"—momentum resamplers, integrators, accept/reject functions, gradient estimators—as pure JAX functions. These atoms are composed using functional combinators, enabling custom, efficient, and parallelizable inference kernels without entanglement in monolithic implementations. Such design empowers modular research, rapid prototyping, and efficient benchmarking.

In probabilistic programming with dynamic computation graphs, combinators (e.g., map, reduce, compose) are used to statically assemble both models and inference evaluation strategies (Sennesh et al., 2018). This enables practitioners to exploit parallelism and conditional independence while retaining black-box Monte Carlo flexibility; the resulting pipelines support training objectives such as variational bounds and wake–sleep optimization.

Composable inference can also appear as metaprogramming over inference itself. By constructing inference metaprograms that dispatch distinct MCMC kernels to independent subproblems and then orchestrate alternation over them with rigorous connectivity conditions (for π-irreducibility, aperiodicity, stationarity), one obtains hybrid samplers with formal convergence guarantees for large and complex probabilistic programs (Handa et al., 2019).

5. Compositionality in Modern Machine Learning, Meta-learning, and Operator Inference

Composable inference is increasingly central in advanced machine learning. In meta-learning, compositional task inference is realized by viewing each task as a structured combination of modules, with adaptation to new tasks performed by probabilistic inference over the composition ("grammar") of modules, without gradient updates (Bakermans et al., 2 Oct 2025). The compositionality of the module grammar enables data-efficient and rapid transfer across tasks.

For operator inference in dynamical systems, structure-preserving neural operators are additively composed, reflecting skew-symmetry, (semi-)positive definiteness, and other algebraic invariants. These operator blocks (e.g., skew, SPSD, unstructured) are combined into a single reduced dynamical model, learned entirely from data, and composability ensures the hybrid model preserves desired physical or geometric properties (Parish et al., 9 Mar 2026).

In energy-based compositional visual inference, concept distributions (EBMs for particular attributes) are combined via sums, LogSumExp, or linear combinations of energies, enabling sampling and inference in conceptually compositional spaces. This logical composability enables zero-shot and continual learning, out-of-distribution generalization, and robust concept inference (Du et al., 2020).

6. Multi-Agent, Multi-Expert, and Distributed Composable Inference

Composable inference frameworks are applied in large-scale, distributed, or multi-agent systems. For example:

• Composite Bayesian inference (Roche, 2015) modularizes the aggregation of views/agents, optimizing the ensemble via convex loss, and balancing interpretability and predictive accuracy.
• In embodied multi-agent systems, composable graphs of thoughts (CGoT) provide the algebra of merging and splitting distributed inference processes (e.g., vehicles and robots), where agents' "thought graphs" are merged/disjointed via explicit operations. This framework is exploited to reduce inference costs (LLM call tokens) without degrading task performance (Nie et al., 25 Oct 2025).
• In LLM inference, Comp-LLM (Sridharan et al., 28 Nov 2025) orchestrates complex reasoning across modular expert LLMs. A query is decomposed into sub-queries, each assigned to a domain expert based on embedding similarity, and a query execution scheduler enforces dependency order and parallelism, culminating in response aggregation by a base LLM. This yields gains in both model size and latency, with no loss in accuracy versus monolithic architectures.

For scalable MAP inference, composable core-sets (e.g., for determinant maximization in DPPs) summarize data partitions independently using algorithms like Greedy or Local Search, and the resulting summaries can be aggregated to yield provably good approximate solutions for the full set—directly leveraging compositionality for optimization scalability (Indyk et al., 2019).

7. Broader Implications, Trade-offs, and Challenges

Composable inference offers several benefits across domains:

• Flexibility and Reusability: New algorithms arise by composing or altering existing components or transformations.
• Interpretability: Modular structure permits inspection, weight learning, and interpretability at the agent or feature level (Roche, 2015).
• Scalability: Distributed or parallel inference (e.g., core-sets, multi-expert, multi-agent coordination) can be achieved with compositional guarantees.
• Correctness and Convergence: Categorical and metaprogramming approaches provide theorem-backed guarantees for asymptotic correctness, modularity, and (under specified connectivities) convergence of hybrid algorithms (Handa et al., 2019, Braithwaite et al., 2023).
• Limits: Interaction of composed effects may be nontrivial when non-local dependencies or global computation graph rewrites are encountered (Moore et al., 2018). Achieving optimal trade-offs between interpretability and predictive performance may depend on weight learning, regularization, or further architectural refinement.

Future work includes advancing compositional effect systems, optimization of handler stacks, algebraic principles for hybrid/inexact compositions, and principled interfaces for parameter and structure learning in modular inference schemes.

---

Major References:

(Roche, 2015, Zinkov et al., 2016, Moore et al., 2018, Sennesh et al., 2018, Braithwaite et al., 2023, Jacobs et al., 28 Nov 2025, Handa et al., 2019, Cabezas et al., 2024, Sridharan et al., 28 Nov 2025, Bakermans et al., 2 Oct 2025, Parish et al., 9 Mar 2026, Du et al., 2020, Nie et al., 25 Oct 2025, Indyk et al., 2019, Smithe, 2021).