SoftWater: Class-Aware Rate Allocation for Softmax Quantization
Di cosa parla
Un nuovo metodo affronta la compressione dello strato finale dei modelli di linguaggio — la parte che trasforma i punteggi in probabilità per scegliere le parole — riducendo lo spazio necessario senza alterare troppo quelle probabilità. La tecnica assegna più precisione alle parole frequenti e meno a quelle rare, basandosi su come le probabilità variano e usando statistiche raccolte con una sola esecuzione per costruire una rappresentazione efficiente. Nei test su vari modelli il metodo supera una soluzione precedente e permette un significativo risparmio di memoria con una piccola perdita nella qualità delle predizioni.
Cosa permette di osservare
Consente di esplorare come bilanciare spazio di memoria e fedeltà delle probabilità finali, come adattare la compressione alla frequenza delle parole e quanto conti scegliere dati di calibrazione coerenti con il dominio d'uso.
Dalla fonte
Post-training quantization pipelines routinely leave the softmax output layer in high precision. Yet in small LLMs with modern vocabularies, the head holds 15--30\% of all parameters, so a nominal ``2-bit'' model with an fp16 head can store several times as many bits per weight. We pose softmax-layer quantization as a rate-distortion problem under the KL divergence between the original and quantized output distributions. A second-order analysis reveals a class-aware geometry: quantization error is weighted jointly by feature covariance and class-specific softmax curvature. A separability approximation replaces the $Kn\times Kn$ Cholesky with one $n\times n$ factorization rescaled per class, making the lattice encodable by successive interference cancellation, with both statistics from a single forward pass. The resulting method, SoftWater, gives fine grids to frequent, low-variance clas…