Accelerating Learned Image Compression Through Modeling Neural Training Dynamics

TL;DR

This paper introduces a novel training acceleration method for learned image compression by modeling neural training dynamics, reducing parameters and variance, and ensuring faster convergence without performance loss.

Contribution

It proposes the STDET and SMA techniques to cluster parameters, embed non-reference parameters efficiently, and smooth training weights, advancing LIC training efficiency.

Findings

Significantly reduces training parameters and space.

Achieves faster convergence without performance loss.

Lower training variance compared to standard SGD.

Abstract

As learned image compression (LIC) methods become increasingly computationally demanding, enhancing their training efficiency is crucial. This paper takes a step forward in accelerating the training of LIC methods by modeling the neural training dynamics. We first propose a Sensitivity-aware True and Dummy Embedding Training mechanism (STDET) that clusters LIC model parameters into few separate modes where parameters are expressed as affine transformations of reference parameters within the same mode. By further utilizing the stable intra-mode correlations throughout training and parameter sensitivities, we gradually embed non-reference parameters, reducing the number of trainable parameters. Additionally, we incorporate a Sampling-then-Moving Average (SMA) technique, interpolating sampled weights from stochastic gradient descent (SGD) training to obtain the moving average weights, ensuring smooth temporal behavior and minimizing training state variances. Overall, our method significantly reduces training space dimensions and the number of trainable parameters without sacrificing model performance, thus accelerating model convergence. We also provide a theoretical analysis on the Noisy quadratic model, showing that the proposed method achieves a lower training variance than standard SGD. Our approach offers valuable insights for further developing efficient training methods for LICs.