Neurosymbolic Diffusion Models

Abstract: Neurosymbolic (NeSy) predictors combine neural perception with symbolic reasoning to solve tasks like visual reasoning. However, standard NeSy predictors assume conditional independence between the symbols they extract, thus limiting their ability to model interactions and uncertainty - often leading to overconfident predictions and poor out-of-distribution generalisation. To overcome the limitations of the independence assumption, we introduce neurosymbolic diffusion models (NeSyDMs), a new class of NeSy predictors that use discrete diffusion to model dependencies between symbols. Our approach reuses the independence assumption from NeSy predictors at each step of the diffusion process, enabling scalable learning while capturing symbol dependencies and uncertainty quantification. Across both synthetic and real-world benchmarks - including high-dimensional visual path planning and rule-based autonomous driving - NeSyDMs achieve state-of-the-art accuracy among NeSy predictors and demonstrate strong calibration.

• The paper introduces NSDMs that use a masked diffusion process to capture inter-concept dependencies, overcoming overconfident predictions in traditional NeSy models.
• The methodology integrates a discrete diffusion model with a symbolic program, employing a variational posterior and a novel NELBO loss to robustly connect unmasked concepts to outputs.
• Experimental results demonstrate improved scalability, reliability, and OOD performance on tasks like MNIST addition and visual path planning.

This paper introduces Neurosymbolic Diffusion Models (NSDMs), a novel class of neurosymbolic (NeSy) predictors that leverage discrete diffusion models to capture dependencies between latent symbols and improve uncertainty quantification. Traditional NeSy predictors often assume conditional independence between symbols extracted by a neural network. This assumption, while enabling efficient probabilistic reasoning, limits their ability to model symbol interactions, leading to overconfident predictions, poor out-of-distribution generalization, and vulnerability to "reasoning shortcuts" (RSs)—where models achieve high output accuracy with incorrect intermediate concepts. NSDMs address this by using a masked diffusion process that operates over the concepts, allowing for global dependencies while retaining local independence at each step of the diffusion process for scalability.

Core Idea and Model Setup

NSDMs integrate a masked diffusion model (MDM) with a symbolic program $\varphi$. The model aims to predict discrete labels $\by \in [V]^Y$ given high-dimensional inputs $\bx$.

1. Concept Extraction (Diffusion Process): Instead of directly predicting concepts $\bc \in [V]^C$, NSDMs define a continuous-time forward diffusion process that gradually masks an initial concept $\bc^0$ to a fully masked state $\bc^1$. A reverse diffusion process, parameterized by a neural network, learns to unmask these concepts.
2. Unmasking Model: The core of the reverse process is a conditionally independent unmasking model $p_\btheta(\tilde{\bc}^0 \mid \bc^t, \bx)$. This model predicts a fully unmasked concept $\tilde{\bc}^0$ given a partially masked concept $\bc^t$ and the input $\bx$. This local independence is crucial for reusing efficient reasoning techniques from standard NeSy.
3. Symbolic Program Integration: The predicted concepts $\by \in [V]^Y$0 are mapped to output labels $\by \in [V]^Y$1 using a given symbolic program $\by \in [V]^Y$2.
4. Latent Concepts: Since ground-truth concepts $\by \in [V]^Y$3 are not available during training (only input-output pairs $\by \in [V]^Y$4), NSDMs employ a variational setup. A variational posterior $\by \in [V]^Y$5 is used to approximate the true posterior $\by \in [V]^Y$6 and provide samples of likely concepts consistent with the observed output.

Principled Loss Function (NELBO)

A key contribution is the derivation of a novel negative evidence lower bound (NELBO) for NSDMs in continuous time:

$\by \in [V]^Y$7

The NELBO consists of three main components:

1. Concept Unmasking Loss ($\by \in [V]^Y$8):

   $\by \in [V]^Y$9

   This term encourages the model to reconstruct the concepts $\bx$0 (sampled from the variational posterior) from their partially masked versions $\bx$1. $\bx$2 is the noising schedule, and $\bx$3 is its derivative.

2. Output Unmasking Loss ($\bx$4):

   $\bx$5

   This loss connects the predicted concepts to the observed outputs via the symbolic program $\bx$6. It's a sum of $\bx$7 weighted model counts (WMCs), one for each output dimension. Because $\bx$8 is conditionally independent, standard NeSy techniques can be used for efficient computation or approximation.

3. Variational Entropy ($\bx$9):

   $\bc \in [V]^C$0

   This term encourages the variational posterior to cover all concepts $\bc \in [V]^C$1 consistent with the input $\bc \in [V]^C$2 and output $\bc \in [V]^C$3, which is crucial for RS-awareness.

To derive this NELBO, the paper proves an extension of the standard MDM NELBO to non-factorized unmasking models, which is necessary because the output reverse process $\bc \in [V]^C$4 does not factorize.

Implementation and Optimization

• Variational Posterior $\bc \in [V]^C$5: Implemented by adapting the MDM's sampling process. It uses the concept unmasking model $\bc \in [V]^C$6 but conditions it on the output $\bc \in [V]^C$7 by re-weighting samples based on consistency with $\bc \in [V]^C$8. $\bc \in [V]^C$9. Since exact sampling is NP-hard, a relaxation is used: $\bc^0$0. Samples are drawn from $\bc^0$1 and re-weighted. In practice, $\bc^0$2 samples are drawn, and the one violating the fewest constraints (or chosen via importance sampling) is selected.
• Loss Optimization:
  • $\bc^0$3: Tractable given samples $\bc^0$4 and $\bc^0$5.
  • $\bc^0$6: Involves WMCs which are #P-hard. For scalability, a REINFORCE Leave-One-Out (RLOO) gradient estimator is used for the WMC terms.
  • $\bc^0$7, where $\bc^0$8 is the empirical mean of constraint satisfaction for output $\bc^0$9 over $\bc^1$0 samples.
  • $\bc^1$1: Maximizing the entropy of the constrained variational posterior is intractable. Two biased approximations are used:
  • 1. Conditional 1-step entropy: $\bc^1$2 (if a tractable circuit for $\bc^1$3 exists).
  • 2. Unconditional 1-step entropy: $\bc^1$4 (otherwise).
  • Hyperparameters $\bc^1$5 weight the loss components. The indirect gradient from sampling $\bc^1$6 is ignored for simplicity and stability.

Architecture Overview

The core neural network component is the unmasking model $\bc^1$7. This network takes the current input $\bc^1$8 and the partially masked concept vector $\bc^1$9 to predict the probabilities for each dimension of the unmasked concept $p_\btheta(\tilde{\bc}^0 \mid \bc^t, \bx)$0. For instance, in visual tasks, $p_\btheta(\tilde{\bc}^0 \mid \bc^t, \bx)$1 could be an image, and $p_\btheta(\tilde{\bc}^0 \mid \bc^t, \bx)$2 could be a grid where some cells have known costs and others are masked. The network would then predict the costs for all cells.

$\bc^t$8

The training process involves:

1. Sampling $p_\btheta(\tilde{\bc}^0 \mid \bc^t, \bx)$3 from data.
2. Sampling $p_\btheta(\tilde{\bc}^0 \mid \bc^t, \bx)$4 (using the re-sampling strategy).
3. Sampling $p_\btheta(\tilde{\bc}^0 \mid \bc^t, \bx)$5 and $p_\btheta(\tilde{\bc}^0 \mid \bc^t, \bx)$6.
4. Computing $p_\btheta(\tilde{\bc}^0 \mid \bc^t, \bx)$7 using $p_\btheta(\tilde{\bc}^0 \mid \bc^t, \bx)$8.
5. Sampling $p_\btheta(\tilde{\bc}^0 \mid \bc^t, \bx)$9.
6. Estimating $\tilde{\bc}^0$0 using RLOO.
7. Estimating $\tilde{\bc}^0$1.
8. Updating $\tilde{\bc}^0$2.

Inference

To predict $\tilde{\bc}^0$3 for a new input $\tilde{\bc}^0$4:

1. Sample $\tilde{\bc}^0$5 concept vectors $\tilde{\bc}^0$6 using the trained MDM (e.g., via first-hitting sampler or T-step discretization).
2. Compute $\tilde{\bc}^0$7 for each sample.
3. Predict $\tilde{\bc}^0$8 (majority voting).

Experiments and Results

NSDMs are evaluated on:

1. Scalability (RQ1):
   • Multidigit MNIST Addition: $\tilde{\bc}^0$9 and $\bc^t$0 digits. NSDM performs competitively with state-of-the-art (SOTA) approximate NeSy methods (A-NeSI) and exact methods, showing its expressiveness doesn't hinder performance on traditional benchmarks.
   • Visual Path Planning: $\bc^t$1 and $\bc^t$2 grids. NSDM significantly outperforms baselines (A-NeSI, EXAL, I-MLE) on the $\bc^t$3 task (concept space $\bc^t$4), demonstrating superior scalability.
2. RS-Awareness (RQ2): Using RSBench (MNIST Half, MNIST Even-Odd, BDD-OIA).
   • Compared to independent models (Semantic Loss, DeepProbLog) and RS-aware ensembles (BEARS).
   • MNIST Tasks: NSDM achieves significantly better concept accuracy (in-distribution and OOD) and much better concept calibration (Expected Calibration Error - ECE), especially with conditional entropy. OOD output accuracy also improves.
   • BDD-OIA (Autonomous Driving): NSDM shows better output F1 than BEARS and significantly improved calibration and concept F1 compared to DeepProbLog.

Practical Implications

• Improved Reliability: By modeling concept dependencies, NSDMs are more aware of ambiguities and reasoning shortcuts, leading to better calibrated predictions. This is crucial for safety-critical applications like autonomous driving.
• Scalability to Complex Problems: The local independence assumption within the diffusion steps and the RLOO-based gradient estimation allow NSDMs to scale to high-dimensional concept spaces where previous expressive NeSy methods struggled.
• Modular Design: The unmasking neural network and the symbolic program are distinct components, allowing flexibility in choosing appropriate architectures and reasoning modules.
• Computational Requirements: Training involves sampling and multiple forward/backward passes for the RLOO estimator, making it more computationally intensive than simpler NeSy models. Inference involves multiple sampling steps from the diffusion model. The symbolic program's complexity also impacts performance.

Limitations and Future Work

• The RLOO estimator relies on decomposing the output $\bc^t$5 into independent dimensions for effective gradient signal.
• Scalability is limited by tasks with efficient circuit representations for $\bc^t$6 or decomposable output spaces.
• The approximations for variational entropy and ignoring the indirect gradient from $\bc^t$7 are simplifications.
• Exploring other discrete diffusion models or hybrid models (discrete/continuous concepts) are future directions.

In summary, Neurosymbolic Diffusion Models offer a promising approach to building more expressive, scalable, and reliable neurosymbolic systems. They achieve this by combining the strengths of discrete diffusion models for capturing dependencies with the efficiency of local independence assumptions and established NeSy reasoning techniques. The empirical results demonstrate SOTA performance in scalability and RS-awareness.