Align Forward, Adapt Backward: Closing the Discretization Gap in Logic Gate Networks
Organizations: Department of Artificial Intelligence and Department of Electrical and Computer Engineering, Inha University, Incheon, Korea.
Abstract
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.
Figures & tables
| Method | =0.05 | =0.1 | =0.5 | =1.0 | =2.0 |
|---|---|---|---|---|---|
| Soft-Mix | 0.0 / 0 | 0.0 / 0 | 0.0 / 0.1 | 0.0 / 0.1 | 0.6 / 24 |
| Soft-Gumbel | 0.7 / 2 | 0.2 / 2 | 0.2 / 7 | 0.2 / 17 | 2.1 / 40 |
| Hard-ST | 0 / 0 | 0 / 0 | 0 / 0 | 0 / 0 | 0 / 0 |
| Gumbel-ST | 40 † | 59 † | 0.0 / 0.1 | 0.0 / 0.3 | 0.0 / 1.1 |
| Hard-ST+ CAGE | 0 / 0 | 0 / 0 | 0 / 0 | 0 / 0 | 0 / 0 |
| Gumbel-ST+ CAGE | 0.0 / 4.5 | 0.1 / 9 | 0.0 / 3.5 | 0.0 / 1.6 | 0.0 / 2.2 |
| Method | MNIST | M-Binary | C-10-Binary |
|---|---|---|---|
| Soft-Mix | 97.1 / 1.4 | 97.2 / 1.1 | 50.7 / 7.5 |
| Soft-Gumbel | 85.9 / 12.0 | 85.6 / 12.2 | 52.6 / 5.4 |
| Hard-ST | 97.1 / 1.3 | 97.1 / 1.1 | 50.9 / 7.2 |
| Gumbel-ST | 52.1 / 46.3 | 51.0 / 47.3 | 18.9 / 39.5 |
| Hard-ST+ CAGE | 98.4 / 0.1 | 98.1 / 0.0 | 57.6 / 0.5 |
| Gumbel-ST+ CAGE | 98.4 / 0.1 | 98.2 / 0.1 | 58.0 / 0.6 |
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
| Domain | Discrete Choices | Training | Inference | Gap? |
|---|---|---|---|---|
| Logic Gate Networks ( Petersen et al., 2022 ; Kim, 2024 ) | Boolean gates | Soft mixture | Hard (argmax) | Yes |
| NAS (DARTS) ( Liu et al., 2019 ) | Operations | Soft mixture | Hard (argmax) | Yes |
| NAS (SNAS/GDAS) ( Xie et al., 2019 ; Dong & Yang, 2019 ) | Operations | Hard (sample) | Hard (argmax) | Small |
| MoE (traditional) ( Jacobs et al., 1991 ) | Expert networks | Soft gating | Soft gating | No |
| MoE (sparse) ( Fedus et al., 2022 ) | Expert networks | Hard (top-k) | Hard (top-k) | No |
| VQ-VAE ( Van Den Oord et al., 2017 ) | Codebook vectors | Hard (nearest) | Hard (nearest) | No |
| Parameter | MNIST | CIFAR-10 |
|---|---|---|
| Neurons per layer | 64,000 | 128,000 |
| Layers | 4, 5, 6 | 4, 5, 6 |
| Gates per node | 16 | 16 |
| Batch size | 512 | 512 |
| Learning rate | 0.01 | 0.05 |
| Optimizer | Adam | Adam |
| MNIST | MNIST-Binary | CIFAR-10-Binary | ||||
|---|---|---|---|---|---|---|
| Method | Peak Gap | Acc Range | Peak Gap | Acc Range | Peak Gap | Acc Range |
| Soft-Mix | 23% | 1.4% | 24% | 1.1% | 9% | 7.5% |
| Soft-Gumbel | 40% | 12.0% | 40% | 12.2% | 3.5% | 5.4% |
| Hard-ST | 0.4% ∗ | 1.3% | 0% | 1.1% | 0% | 7.2% |
| Gumbel-ST | 62% † | 46.3% | 61% † | 47.3% | 16% † | 39.5% |
| Hard-ST+CAGE | 0.3% ∗ | 0.1% | 0% | 0.0% | 0% | 0.5% |
| Method | L | =0.05 | =0.1 | =0.5 | =1.0 | =2.0 |
|---|---|---|---|---|---|---|
| Soft-Mix | 4 | [+0.2,+0.4] | [+0.1,+0.3] | [+0.1,+0.3] | [+0.0,+0.3] | [+0.2,+0.5] |
| 5 | [+0.2,+0.4] | [+0.1,+0.3] | [+0.1,+0.3] | [+0.1,+0.3] | [+0.2,+1.5] | |
| 6 | [+0.1,+0.4] | [+0.1,+0.3] | [+0.1,+0.3] | [+0.1,+0.3] | [+0,+23] | |
| Soft-Gumbel | 4 | [-31,+2] | [-15,+2] | [-0.1,+2.3] | [+0,+10] | [+0,+21] |
| 5 | [-33,+1] | [-33,+2] | [+0.0,+2.9] | [+0,+16] | [+1,+33] | |
| 6 | [-29,+4] | [-25,+4] | [-0,+9] | [+0,+16] | [+1,+40] |
| Method | L | =0.05 | =0.1 | =0.5 | =1.0 | =2.0 |
|---|---|---|---|---|---|---|
| Soft-Mix | 4 | [-0.1,+0.0] | [-0.1,+0.1] | [-0.1,+0.2] | [-0.1,+0.3] | |
| 5 | [-0.1,+0.1] | [-0.1,+0.1] | [-0.1,+1.6] | |||
| 6 | [-0.1,+0.1] | [-0.1,+0.1] | [+0,+24] | |||
| Soft-Gumbel | 4 | [-31,+1] | [-14,+1] | [-0.3,+1.8] | [+0,+9] | [+0,+23] |
| 5 | [-35,+1] | [-31,+2] | [-0.4,+4.3] | [+0,+18] | [+0,+35] | |
| 6 | [-30,+2] | [-25,+2] | [-1,+7] | [+0,+17] | [+2,+40] |
| Method | L | =0.05 | =0.1 | =0.5 | =1.0 | =2.0 |
|---|---|---|---|---|---|---|
| Soft-Mix | 4 | [-0.1,+0.1] | [-0.1,+0.1] | [-0.1,+0.1] | [-0.3,+0.3] | [-0.3,+0.5] |
| 5 | [-0.1,+0.1] | [-0.1,+0.1] | [-0.1,+0.1] | [-0.3,+0.3] | [-0.2,+1.5] | |
| 6 | [-0.1,+0.1] | [-0.1,+0.1] | [-0.1,+0.1] | [-0.2,+0.3] | [-0,+9] | |
| Soft-Gumbel | 4 | [-2.3,+0.4] | [-2.0,+0.4] | [-0.7,+0.5] | [-0.3,+0.3] | [-0.2,+0.4] |
| 5 | [-2.5,+0.3] | [-2.6,+0.4] | [-0.8,+0.7] | [-0.3,+0.5] | [-0.3,+0.4] | |
| 6 | [-3.5,+0.2] | [-2.3,+0.4] | [-0.7,+0.5] | [-0.3,+0.5] | [-0.3,+0.4] |
| Claim | Comparison | Stat. | Prac. | |||
|---|---|---|---|---|---|---|
| Temperature sensitivity (accuracy) | ||||||
| Gumbel-ST fails at low | =0.05 vs =2.0 | 36.16 | .001 | 31.21 | ✓ | |
| Hard-ST is -robust | =0.05 vs =2.0 | 14.53 | .005 | 16.26 | ** | – |
| CAGE effectiveness | ||||||
| CAGE fixes Gumbel-ST | Gumbel-ST vs +CAGE ( =0.05) | 38.81 | .001 | 31.28 | ✓ | |
| Selection gap comparisons (at =2.0) | ||||||
| Claim | MNIST | MNIST-Binary | CIFAR-10-Binary |
|---|---|---|---|
| Gumbel-ST fails at low | |||
| Hard-ST is -robust | ** † | ** † | *** † |
| CAGE fixes Gumbel-ST |
| Index | Gate | Boolean | Soft Formula | Hard Formula |
|---|---|---|---|---|
| 0 | FALSE | |||
| 1 | AND | |||
| 2 | A B | |||
| 3 | A | |||
| 4 | A B | |||
| 5 | B |
| Distribution | Unary | AND/OR | XOR |
|---|---|---|---|
| Binary (lower bound) | 0% | 0% | 0% |
| MNIST (91% binary-like) | 2% | 3% | 5% |
| Uniform | 25% | 21.9% | 16.7% |
| At boundary (upper) | 50% | 75% | 50% |
| Category | Gates | Gap (%) |
|---|---|---|
| Constants | FALSE, TRUE | 0.0 |
| XOR-type | XOR, XNOR | 16.7 |
| AND/OR-type | AND, OR, NAND, NOR, A B, A B, A B, A B | 21.9 |
| Unary | A, B, A, B | 25.0 |
| Pixel Range | MNIST [0,1] | MNIST-Binary |
|---|---|---|
| Exactly 0 | 80.9% | 86.7% |
| 1.7% | 0.0% | |
| 8.7% | 0.0% | |
| 8.8% | 13.3% | |
| Binary-like ( or ) | 91.3% | 100.0% |
Neural Network Optimization
Neural network optimization studies algorithms and training dynamics that help models learn efficiently, reliably, and generalize well.
See all 321 papers