Title: Implement Tractable Dropout Inference (TDI) across all SPFlow circuit modules (Ventola et al., UAI 2023)

Summary

Implement tractable dropout inference (TDI) in SPFlow as a general inference-time routine that provides sampling-free epistemic uncertainty estimates (at least expectation + variance, optionally covariance later) for all SPFlow probabilistic circuit modules, not just RAT‑SPNs. The goal is to enable uncertainty signals (e.g., predictive entropy / variance-based scores) for OOD detection and robustness under distribution shift in a single forward pass.

Motivation

Ventola et al. show that conventional probabilistic circuits (PCs) can be overconfident under OOD and distribution shift. They propose TDI, a closed-form approximation to Monte Carlo dropout for PCs using variance propagation, yielding uncertainty estimates without multiple stochastic forward passes.

References

• Reference implementation repo: https://github.com/ml-research/tractable-dropout-inference

Scope

In scope

• A generic TDI inference path that works for arbitrary SPFlow circuit graphs composed of existing SPFlow modules, including:

  • Sum-family: Sum, ElementwiseSum, SignedSum, RepetitionMixingLayer, convolutional sums (e.g., SumConv)
  • Product-family: Product, ElementwiseProduct, OuterProduct, convolutional products (e.g., ProdConv), Einsum-style combinators if applicable
  • Structure/ops: Cat, split modules, factorization modules, wrappers
  • Leaves: all leaf distributions in spflow/modules/leaves/
  • Zoo models: RatSPN, Einet, and other bundled architectures should benefit automatically if they are built from supported modules

• Output should include (where meaningful):

  • point predictions (mean / expected log-likelihood or consistent equivalent)
  • uncertainty proxy (variance; optionally derived scalar metrics)
  • classification-facing uncertainty: predictive entropy and/or posterior variance

• Provide a Monte Carlo dropout (MCD) baseline for validation (not necessarily as a public API, but usable in tests/benchmarks).

Out of scope (initially)

• Full covariance propagation for arbitrary DAGs if it makes the system impractical (time/memory); however, the implementation must define a clear policy for covariance handling.
• End-to-end reproduction of all paper experiments (datasets, corruptions, plots, hyperparameter sweeps), unless separately requested.

Key Requirements / Design Decisions

1. Dropout semantics (global policy)

• Define how dropout applies to circuit edges/nodes across module types:
  • Prefer a consistent interpretation of “dropout in PCs” that maps onto tensorized modules (activation-style dropout at mixture inputs is likely the most implementable).
  • Clarify whether dropout is applied only to sum/mixture operations or also to other module families.

2. Moment convention

• Decide whether to propagate:
  • moments of likelihood in linear space vs log-space representations,
  • and ensure consistency with SPFlow’s log_likelihood API and existing caching.

3. Generic propagation API

• Define a generic “TDI forward” contract for modules: each module must be able to consume child moments and output its own:
  • E (expectation / mean signal)
  • Var (variance)
  • (optional) Cov or covariance handling strategy

4. Covariance handling strategy

• Must be explicitly defined:
  • Default: covariance-free approximation (“copy-paste”/ignore covariance terms) for tractability
  • Optional extensions:
    • structure-aware zero covariance (trees / provably independent subgraphs)
    • covariance bounds (Cauchy-Schwarz) where applicable
    • full covariance propagation only for restricted structures

5. Classification support

• Provide a standardized way to compute posterior uncertainty for multi-class PCs:
  • moment approximation for ratios (Taylor expansion) to produce posterior mean/variance
  • predictive entropy as a first-class output for OOD thresholding

6. Numerical stability

• Ensure the propagation is stable across module compositions:
  • handle subtraction/cancellation in variance formulas
  • define a policy for tiny negative variances from numerical error (clamp, log-domain safe ops, etc.)
  • avoid NaNs/Infs for typical configurations

Deliverables / Acceptance Criteria

• A general TDI inference entrypoint that works on circuits built from supported SPFlow modules and returns:
  • expectation-like output compatible with existing inference semantics
  • variance (or log-variance) and derived uncertainty metrics
• Coverage across all core module families (sum/product/ops/leaves/conv/einsum where applicable)
• Tests:
  • MCD vs TDI moment matching on small circuits for multiple module types (not only RAT)
  • shape/broadcasting invariants across repetitions/channels/features
  • numerical stability tests (no NaNs; variance policy enforced)
• Minimal documentation describing:
  • dropout parameter semantics and where it applies
  • supported modules and known limitations
  • interpretation of returned quantities

Proposed Milestones

M1 — Core framework + baseline modules

• Establish propagation API and default covariance policy
• Implement TDI for foundational modules: Sum, Product, leaf base + common leaves, Cat, basic splits

M2 — Broaden module coverage

• Add support for remaining sum/product variants (ElementwiseSum, SignedSum, ElementwiseProduct, OuterProduct, etc.)
• Ensure zoo architectures (RAT‑SPN, Einet) work end-to-end without special casing

M3 — Specialized layers

• Convolutional modules (SumConv, ProdConv), einsum/linsum layers, wrappers, any remaining ops

M4 — Validation + stabilization

• Expand MCD-vs-TDI tests across module zoo
• Benchmark performance overhead
• Tighten numerical stability and variance handling

Notes

This issue defines the conceptual scope and acceptance criteria for a full-SPFlow implementation of TDI. No direct code changes are proposed in this issue description.