# Align Forward, Adapt Backward: Closing the Discretization Gap in Logic Gate Networks

arXiv: https://arxiv.org/abs/2603.14157 · HTML: https://mlsanity.com/paper/2603.14157/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 37 references that are in the map.
- First posted on arXiv 14 Mar 2026.

## Builds on (most cited first, 24 of 24)

- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [SNAS: Stochastic Neural Architecture Search](https://mlsanity.com/paper/1812.09926/) (2018)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Searching for A Robust Neural Architecture in Four GPU Hours](https://mlsanity.com/paper/1910.04465/) (2019)
- [From Sparse to Soft Mixtures of Experts](https://mlsanity.com/paper/2308.00951/) (2023)
- [A Survey on Mixture of Experts in Large Language Models](https://mlsanity.com/paper/2407.06204/) (2024)
- [Progressive Differentiable Architecture Search: Bridging the Depth Gap between Search and Evaluation](https://mlsanity.com/paper/1904.12760/) (2019)
- [Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints](https://mlsanity.com/paper/2212.05055/) (2022)
- [Understanding and Robustifying Differentiable Architecture Search](https://mlsanity.com/paper/1909.09656/) (2019)
- [MoEfication: Transformer Feed-forward Layers are Mixtures of Experts](https://mlsanity.com/paper/2110.01786/) (2021)
- [Fair DARTS: Eliminating Unfair Advantages in Differentiable Architecture Search](https://mlsanity.com/paper/1911.12126/) (2019)
- [DARTS+: Improved Differentiable Architecture Search with Early Stopping](https://mlsanity.com/paper/1909.06035/) (2019)
- [Analyzing Differentiable Fuzzy Logic Operators](https://mlsanity.com/paper/2002.06100/) (2020)
- [Deep Differentiable Logic Gate Networks](https://mlsanity.com/paper/2210.08277/) (2022)
- [Efficient Automation of Neural Network Design: A Survey on Differentiable Neural Architecture Search](https://mlsanity.com/paper/2304.05405/) (2023)
- [Towards Narrowing the Generalization Gap in Deep Boolean Networks](https://mlsanity.com/paper/2409.05905/) (2024)
- [Small Temperature is All You Need for Differentiable Architecture Search](https://mlsanity.com/paper/2306.06855/) (2023)

## Built on by (most cited first, 2 of 2)

- [Fitting Multilinear Polynomials for Logic Gate Networks](https://mlsanity.com/paper/2605.08657/) (2026)
- [Task-Aware Discretization of Differentiable Logic Gate Networks](https://mlsanity.com/paper/2609.33747/) (2026)

## Abstract (arXiv)

In neural network models, soft mixtures of fixed candidate components (e.g., logic gates and sub-networks) are often used during training for stable optimization, while hard selection is typically used at inference. This raises questions about training-inference mismatch. We analyze this gap by separating forward-pass computation (hard selection vs. soft mixture) from stochasticity (with vs. without Gumbel noise). Using logic gate networks as a testbed, we observe distinct behaviors across four methods: Hard-ST achieves zero selection gap by construction; Gumbel-ST achieves near-zero gap when training succeeds but suffers accuracy collapse at low temperatures; Soft-Mix achieves small gap only at low temperature via weight concentration; and Soft-Gumbel exhibits large gaps despite Gumbel noise, confirming that noise alone does not reduce the gap. We propose CAGE (Confidence-Adaptive Gradient Estimation) to maintain gradient flow while preserving forward alignment. On logic gate networks, Hard-ST with CAGE achieves over 98% accuracy on MNIST and over 58% on CIFAR-10, both with zero selection gap across all temperatures, while Gumbel-ST without CAGE suffers a 47-point accuracy collapse.