Strengthening learnable quantizers for quantization-aware training in low-bit precision

TinyML focuses on deploying Neural Networks on memory-constrained devices, requiring effective model compression. Quantization, which represents weights and activations with low-bit integers, is widely adopted. The traditional approach is Post-Training Quantization (PTQ), which often causes significant accuracy loss. Quantization-Aware Training (QAT) mitigates this issue by fine-tuning the network while including quantization operations, enabling the model to adapt to discrete elements. However, QAT alone is still insufficient at very low bitwidths. Among the techniques that have been proposed to improve QAT, we focus on learnable quantizers, where quantization parameters, in our case scale and offset, are optimized with model weights. Since the distribution of weights and activations in a network can vary significantly during model training, allowing the quantization functions to adapt accordingly appears to be a promising approach. We propose a customizable framework that strengthens learnable quantizers technique, especially in low-bit quantization conditions, through three novel contributions: (i) enforcement of quantization parameters constraints through the RoundClipZero procedure, to ensure offset validity during training; (ii) a tailored gradient rescaling, which improves the effectiveness of offset updates; and (iii) a novel training strategy, consisting in an alternated optimization scheme for model and quantization parameters. Image Classification experiments demonstrate that our framework recovers full-precision accuracy in all cases. Our ablation study of the gradient rescaling and alternated training strategy shows that the performance improvement is generally negligible, but in challenging settings it could make a difference. Mixed-precision experiments reveal insights into layer-wise sensitivity and bit allocation strategies. Comparison with two existing compression frameworks confirms the superiority of our approach, especially in critical quantization conditions. Finally, our results highlight that carefully handling constraints, gradients, and training in learnable quantizers is essential to achieve an accurate quantized model, ready for on-device inference.

Il TinyML si concentra sul deployment di Reti Neurali su dispositivi con vincoli di memoria e richiede un'efficiente compressione dei modelli. La quantizzazione, che rappresenta pesi e attivazioni con interi a pochi bit, è molto diffusa. L'approccio tradizionale è la Post-Training Quantization (PTQ), che spesso causa una perdita significativa di accuratezza. Il Quantization-Aware Training (QAT) mitiga questo problema con un fine-tuning che include le operazioni di quantizzazione, adattando il modello agli elementi discreti, ma rimane insufficiente a precisioni molto basse. Tra le tecniche che sono state proposte per migliorare il QAT, ci concentriamo sulla quella dei learnable quantizers, dove i parametri di quantizzazione, nel nostro caso scale e offset, sono ottimizzati con i parametri del modello. Poiché la distribuzione di pesi e attivazioni varia durante l’allenamento, è promettente far adattare le funzioni di quantizzazione di conseguenza. Proponiamo un framework personalizzabile che consolida la tecnica dei learnable quantizers, specialmente in condizioni di quantizzazione a pochi bit, tramite tre contributi innovativi: (i) imposizione dei vincoli dei parametri di quantizzazione tramite la procedura RoundClipZero per mantenere la validità degli offset durante l'allenamento; (ii) un rescaling dei gradienti ad hoc, che migliora l'efficacia dell'allenamento degli offset; (iii) una strategia di allenamento innovativa, che consiste in uno schema di ottimizzazione alternato per parametri del modello e di quantizzazione. Gli esperimenti di classificazione di immagini dimostrano che il nostro framework recupera l'accuratezza del modello a piena precisione in tutti i casi. Il nostro studio di ablazione, condotto sul rescaling dei gradienti e sulla strategia di allenamento alternato, mostra che il miglioramento delle prestazioni è generalmente trascurabile, ma potrebbe fare la differenza in scenari critici. Gli esperimenti con precisione mista rivelano intuizioni sulla sensibilità dei layer e sulle strategie di allocazione dei bit. Il confronto con due framework di compressione esistenti conferma la superiorità del nostro approccio, soprattutto in condizioni di quantizzazione critiche. Infine, i nostri risultati evidenziano che gestire attentamente vincoli, gradienti e allenamento nei learnable quantizer è essenziale per ottenere un modello quantizzato accurato, pronto per l'inferenza sul dispositivo.