Generalized and Optimal Straight-Through Estimators

James Hooper, Alexander Shekhovtsov
Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, PMLR 300:3187-3195, 2026.

Abstract

Modern ML models often utilize discrete components within their computational graphs, making training challenging. In such cases, approximate-chain-rule gradient estimators can be applied. They work reasonably well but are obtained by combining diverse rationales with ad-hoc choices. In this work, we propose a principled axiomatic approach to define a general family of gradient estimators and show that it subsumes many existing methods. Within this family, we derive optimal estimators with respect to a minimum variance criterion subject to interpretable bias-limiting constraints, addressing integer and one-hot categorical discrete variables. We empirically demonstrate that our estimator can achieve a better bias-variance trade-off than existing ones on synthetic problems and outperforms them on training variational auto-encoders with discrete latent variables.

Cite this Paper

BibTeX
@InProceedings{pmlr-v300-hooper26a, title = { Generalized and Optimal Straight-Through Estimators }, author = {Hooper, James and Shekhovtsov, Alexander}, booktitle = {Proceedings of The 29th International Conference on Artificial Intelligence and Statistics}, pages = {3187--3195}, year = {2026}, editor = {Khan, Emtiyaz and Li, Yingzhen and Solin, Arno and Ramdas, Aaditya}, volume = {300}, series = {Proceedings of Machine Learning Research}, month = {02--05 May}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v300/main/assets/hooper26a/hooper26a.pdf}, url = {https://proceedings.mlr.press/v300/hooper26a.html}, abstract = { Modern ML models often utilize discrete components within their computational graphs, making training challenging. In such cases, approximate-chain-rule gradient estimators can be applied. They work reasonably well but are obtained by combining diverse rationales with ad-hoc choices. In this work, we propose a principled axiomatic approach to define a general family of gradient estimators and show that it subsumes many existing methods. Within this family, we derive optimal estimators with respect to a minimum variance criterion subject to interpretable bias-limiting constraints, addressing integer and one-hot categorical discrete variables. We empirically demonstrate that our estimator can achieve a better bias-variance trade-off than existing ones on synthetic problems and outperforms them on training variational auto-encoders with discrete latent variables. } }
Endnote
%0 Conference Paper %T Generalized and Optimal Straight-Through Estimators %A James Hooper %A Alexander Shekhovtsov %B Proceedings of The 29th International Conference on Artificial Intelligence and Statistics %C Proceedings of Machine Learning Research %D 2026 %E Emtiyaz Khan %E Yingzhen Li %E Arno Solin %E Aaditya Ramdas %F pmlr-v300-hooper26a %I PMLR %P 3187--3195 %U https://proceedings.mlr.press/v300/hooper26a.html %V 300 %X Modern ML models often utilize discrete components within their computational graphs, making training challenging. In such cases, approximate-chain-rule gradient estimators can be applied. They work reasonably well but are obtained by combining diverse rationales with ad-hoc choices. In this work, we propose a principled axiomatic approach to define a general family of gradient estimators and show that it subsumes many existing methods. Within this family, we derive optimal estimators with respect to a minimum variance criterion subject to interpretable bias-limiting constraints, addressing integer and one-hot categorical discrete variables. We empirically demonstrate that our estimator can achieve a better bias-variance trade-off than existing ones on synthetic problems and outperforms them on training variational auto-encoders with discrete latent variables.
APA
Hooper, J. & Shekhovtsov, A.. (2026). Generalized and Optimal Straight-Through Estimators . Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research 300:3187-3195 Available from https://proceedings.mlr.press/v300/hooper26a.html.

Related Material