[Paper Review] Injecting Logical Constraints into Neural Networks via Straight-Through Estimators

Zhun Yang, Joohyung Lee+1 more|arXiv (Cornell University)|Jul 10, 2023

Neural Networks and Applications4 citations

TL;DR

This paper proposes CL-STE, a method that injects discrete logical constraints into neural networks using straight-through estimators (STE) to enable differentiable training. By encoding propositional logic constraints as a differentiable loss function and leveraging STE for gradient backpropagation, the approach achieves scalable, GPU-accelerated neuro-symbolic learning that outperforms symbolic methods in training speed while maintaining accuracy on tasks like shortest path prediction and semi-supervised classification.

ABSTRACT

Injecting discrete logical constraints into neural network learning is one of the main challenges in neuro-symbolic AI. We find that a straight-through-estimator, a method introduced to train binary neural networks, could effectively be applied to incorporate logical constraints into neural network learning. More specifically, we design a systematic way to represent discrete logical constraints as a loss function; minimizing this loss using gradient descent via a straight-through-estimator updates the neural network's weights in the direction that the binarized outputs satisfy the logical constraints. The experimental results show that by leveraging GPUs and batch training, this method scales significantly better than existing neuro-symbolic methods that require heavy symbolic computation for computing gradients. Also, we demonstrate that our method applies to different types of neural networks, such as MLP, CNN, and GNN, making them learn with no or fewer labeled data by learning directly from known constraints.

Motivation & Objective

• To address the challenge of incorporating discrete logical constraints into neural network training, which is difficult due to non-differentiability of logical operations.
• To overcome the computational bottleneck of existing neuro-symbolic methods that rely on expensive symbolic computation (e.g., weighted model counting or solver calls) for gradient computation.
• To develop a scalable, differentiable framework that enables end-to-end training of neural networks with explicit logical constraints using standard deep learning pipelines.
• To demonstrate the effectiveness of CL-STE across diverse architectures (MLP, CNN, GNN) and learning paradigms (supervised, semi-supervised) with minimal labeled data.

Proposed method

• The method encodes logical constraints in conjunctive normal form (CNF) as a differentiable loss function $ L_{cnf} $, which measures the violation of logical clauses by the network's output.
• It applies a straight-through estimator (STE) to the binarized output of the network, allowing gradient flow during backpropagation despite the non-differentiable binarization function.
• The binarization function $ b(x) $ maps real-valued logits to binary values (0 or 1), and the STE replaces the zero gradient of $ b(x) $ with the gradient of a continuous relaxation during backpropagation.
• The total loss combines the logical constraint loss $ L_{cnf} $ and a standard cross-entropy loss $ L_{cross} $, with a hyperparameter $ \alpha $ balancing the two.
• The approach uses a training gate function (TGF) as a differentiable proxy for the binarization process, enabling gradient updates that steer the network toward satisfying logical constraints.
• The method is compatible with various neural architectures (MLP, CNN, GNN) and supports batch training on GPUs, significantly improving training efficiency over symbolic alternatives.

Experimental results

Research questions

• RQ1 Can logical constraints expressed in propositional logic be effectively encoded as a differentiable loss function for neural network training?
• RQ2 Can the straight-through estimator (STE) be effectively adapted to enable gradient-based optimization of neural networks under discrete logical constraints?
• RQ3 Does CL-STE achieve comparable or better performance than existing neuro-symbolic methods while significantly reducing training time through GPU-accelerated, differentiable training?
• RQ4 Can CL-STE generalize across diverse neural network architectures (MLP, CNN, GNN) and learning settings (supervised, semi-supervised)?

Key findings

• CL-STE achieved comparable test accuracy to NeurASP on the shortest path problem in 179 seconds, while NeurASP required significantly longer training time due to symbolic computation.
• On the MNIST and Fashion-MNIST datasets, CL-STE achieved similar classification accuracy to the semantic loss method from Xu et al. (2018), with a training time of ~10 minutes compared to ~12 minutes for the original method.
• The method reduced training time by a factor of ~10 compared to the semantic loss method when trained on CPU, and remained efficient even with complex constraints.
• In the shortest path problem, CL-STE reached 83% accuracy in 83 seconds (baseline) and 85% in 179 seconds (CL-STE), outperforming the semantic loss method, which required 1032 seconds to reach 30.75% accuracy on CPU.
• The method demonstrated scalability and efficiency by leveraging batch training and GPU acceleration, avoiding the computational overhead of constructing arithmetic circuits or calling symbolic solvers.