Binary and ternary AI models train faster with less memory on edge devices
Noise-shaped Sampling and Annealing for Scalable On-the-edge Training of Binary and Ternary AI Models
Machine Learning
Summary
Training AI models that use only a few values for their connections (binary or ternary networks) is challenging because it’s hard to find the best settings that work well on new data. The authors linked this challenge to a type of physics problem called an Ising model and designed a new training method called the ΣΔ-Sampler. This method cleverly adds noise and smooths out the search process, helping the model learn effectively while using less memory and running faster. They tested it on various AI tasks, showing it almost matches the accuracy of full-detail models but in a much more efficient way.
What this means in practice
- •For edge ai developers: Train large binary and ternary AI models on edge devices with less memory and faster updates.
- •For machine learning engineers: Fine-tune discrete-weight models on-device efficiently while preserving accuracy and reducing compute demands.
Authors
Zihao Chen, Zhili Xiao, Rui-Jie Zhu, Yipeng Wang, Jason Eshraghian, Shantanu Chakrabartty
Abstract
The training loss of binary and ternary neural networks, expanded as a multilinear polynomial over the discrete weights, is structurally equivalent to the energy of a hierarchical high-order Ising model. We exploit this equivalence to apply Ising solvers to discrete-weight network training, while accounting for a fundamental difference in objectives: an Ising solver seeks the ground state of a fixed energy function, whereas training seeks solutions that generalize to unseen data. Our key contribution is $ΣΔ$-Sampler, an optimizer that reconciles these objectives by embedding a stochastic sampling-and-annealing step as a noisy quantizer inside a first-order $ΣΔ$ noise-shaping loop. The loop pushes noise from weight quantization, high-order frustration, and mini-batch sampling out of the low-frequency adaptation band, so that discrete weight updates follow the smooth descent trajectories embedded in the rugged loss landscape. We validate the framework on a planted quadratic unconstrained binary optimization (QUBO) benchmark with a known minimizer; on CIFAR-10, where a $13$M-parameter ternary model comes within $1\%$ of full-precision accuracy; on ImageNet-1k, where a ternary ResNet-18 comes within $0.4\%$ Top-5 accuracy of a ternary baseline; on class-incremental continual learning, where prior-task accuracy is preserved; and on MatMul-free language models up to $2.7$B parameters, where it outperforms a straight-through-estimator baseline under matched continued pre-training. Compared with standard quantization-aware training, $ΣΔ$-Sampler reduces the training memory footprint by over $50\%$, delivers an order-of-magnitude speed-up, and reduces weight-update memory traffic by four orders of magnitude, making it attractive for on-device training and fine-tuning of large binary and ternary models.