Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks

要約

本研究では、ベクトル量子化を用いたニューラルネットワークをストレートスルー推定でトレーニングする際の課題を検討する。学習が不安定になる主な原因は、モデルの埋め込みとコードベクトル分布の不一致であることを発見した。この問題の要因として、コードブック勾配のスパース性、コミットメントロスの非対称性などを特定し、コードベクトル割り当てのズレを引き起こすとした。我々は、コードベクトルのアフィン再パラメータ化によってこの問題を解決することを提案する。さらに、ストレートスルー推定によってもたらされる勾配誤差を低減するために、交互最適化を導入する。さらに、コードブック表現とモデル埋め込みとの間のより良い整合を確保するために、コミットメントロスの改良を提案する。これらの最適化手法は、ストレートスルー推定の数学的近似を改善し、最終的にモデル性能を向上させる。我々は、画像分類や生成モデリングを含む様々なタスクにおいて、AlexNet、ResNet、ViTなど、いくつかの一般的なモデルアーキテクチャで、我々の方法の有効性を実証する。

要約(オリジナル)

This work examines the challenges of training neural networks using vector quantization using straight-through estimation. We find that a primary cause of training instability is the discrepancy between the model embedding and the code-vector distribution. We identify the factors that contribute to this issue, including the codebook gradient sparsity and the asymmetric nature of the commitment loss, which leads to misaligned code-vector assignments. We propose to address this issue via affine re-parameterization of the code vectors. Additionally, we introduce an alternating optimization to reduce the gradient error introduced by the straight-through estimation. Moreover, we propose an improvement to the commitment loss to ensure better alignment between the codebook representation and the model embedding. These optimization methods improve the mathematical approximation of the straight-through estimation and, ultimately, the model performance. We demonstrate the effectiveness of our methods on several common model architectures, such as AlexNet, ResNet, and ViT, across various tasks, including image classification and generative modeling.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL