Layer-wise Quantization for Quantized Optimistic Dual Averaging

TL;DR

This paper introduces a layer-wise quantization framework and a novel QODA algorithm for distributed variational inequalities, improving training efficiency of deep neural networks with heterogenous layers.

Contribution

It proposes a general layer-wise quantization method with tight bounds and a new QODA algorithm with adaptive learning rates for monotone variational inequalities.

Findings

QODA achieves up to 150% speedup in training Wasserstein GANs.

Layer-wise quantization adapts to heterogeneity across neural network layers.

The framework provides tight variance and code-length bounds.

Abstract

Modern deep neural networks exhibit heterogeneity across numerous layers of various types such as residuals, multi-head attention, etc., due to varying structures (dimensions, activation functions, etc.), distinct representation characteristics, which impact predictions. We develop a general layer-wise quantization framework with tight variance and code-length bounds, adapting to the heterogeneities over the course of training. We then apply a new layer-wise quantization technique within distributed variational inequalities (VIs), proposing a novel Quantized Optimistic Dual Averaging (QODA) algorithm with adaptive learning rates, which achieves competitive convergence rates for monotone VIs. We empirically show that QODA achieves up to a $150\%$ speedup over the baselines in end-to-end training time for training Wasserstein GAN on $12+$ GPUs.