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

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 and Gumbel-ST with CAGE both achieve over 98\% accuracy on MNIST and ${\sim}58\%$ on CIFAR-10; Hard-ST with CAGE further maintains zero selection gap across all temperatures by construction, while Gumbel-ST without CAGE suffers a ${\sim}47$-point accuracy collapse.

Imagine painting a picture, but in practice you must blend all your colors into one muddy mix and are judged only on that blend. Then the contest changes the rule: you may use just one pure color, no mixing. Because you only practiced the blend, you don't know which pure color to use, and the painting comes out worse. Many neural-network AI models face the same problem, like learnable logic circuits or architecture search that must pick one operation or one subnetwork at each step. Such discrete choices can't be trained directly, so they practice by blending every option, then commit one at deployment. That switch, from blended practice to a single committed choice, quietly lowers accuracy. People thought the fix was a dash of randomness that keeps such choices trainable; we found the real key is simpler. Just practice the same single choices you will use for real, a principle we call forward alignment. The committed choice stays in the foreground, while the gentle learning signal and a built-in tuner work in the background to keep training from collapsing. The idea fits any AI model that must commit to discrete choices, letting it deploy without the usual accuracy drop.