Custom Gradient Estimators are Straight-Through Estimators in Disguise

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Schoenbauer, Matt, Moro, Daniele, Lew, Lukasz, Howard, Andrew
Formato: Preprint
Publicado: 2024
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866929353423585280
author Schoenbauer, Matt
Moro, Daniele
Lew, Lukasz
Howard, Andrew
author_facet Schoenbauer, Matt
Moro, Daniele
Lew, Lukasz
Howard, Andrew
contents Quantization-aware training comes with a fundamental challenge: the derivative of quantization functions such as rounding are zero almost everywhere and nonexistent elsewhere. Various differentiable approximations of quantization functions have been proposed to address this issue. In this paper, we prove that when the learning rate is sufficiently small, a large class of weight gradient estimators is equivalent with the straight through estimator (STE). Specifically, after swapping in the STE and adjusting both the weight initialization and the learning rate in SGD, the model will train in almost exactly the same way as it did with the original gradient estimator. Moreover, we show that for adaptive learning rate algorithms like Adam, the same result can be seen without any modifications to the weight initialization and learning rate. We experimentally show that these results hold for both a small convolutional model trained on the MNIST dataset and for a ResNet50 model trained on ImageNet.
format Preprint
id arxiv_https___arxiv_org_abs_2405_05171
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Custom Gradient Estimators are Straight-Through Estimators in Disguise
Schoenbauer, Matt
Moro, Daniele
Lew, Lukasz
Howard, Andrew
Machine Learning
Quantization-aware training comes with a fundamental challenge: the derivative of quantization functions such as rounding are zero almost everywhere and nonexistent elsewhere. Various differentiable approximations of quantization functions have been proposed to address this issue. In this paper, we prove that when the learning rate is sufficiently small, a large class of weight gradient estimators is equivalent with the straight through estimator (STE). Specifically, after swapping in the STE and adjusting both the weight initialization and the learning rate in SGD, the model will train in almost exactly the same way as it did with the original gradient estimator. Moreover, we show that for adaptive learning rate algorithms like Adam, the same result can be seen without any modifications to the weight initialization and learning rate. We experimentally show that these results hold for both a small convolutional model trained on the MNIST dataset and for a ResNet50 model trained on ImageNet.
title Custom Gradient Estimators are Straight-Through Estimators in Disguise
topic Machine Learning
url https://arxiv.org/abs/2405.05171