2025-11-13T04:07:09.837900

Optimal Quantization for Matrix Multiplication

Ordentlich, Polyanskiy
Recent work in machine learning community proposed multiple methods for performing lossy compression (quantization) of large matrices. This quantization is important for accelerating matrix multiplication (main component of large language models), which is often bottlenecked by the speed of loading these matrices from memory. Unlike classical vector quantization and rate-distortion theory, the goal of these new compression algorithms is to be able to approximate not the matrices themselves, but their matrix product. Specifically, given a pair of real matrices $A,B$ an encoder (compressor) is applied to each of them independently producing descriptions with $R$ bits per entry. These representations subsequently are used by the decoder to estimate matrix product $A^\top B$. In this work, we provide a non-asymptotic lower bound on the mean squared error of this approximation (as a function of rate $R$) for the case of matrices $A,B$ with iid Gaussian entries. Algorithmically, we construct a universal quantizer based on nested lattices with an explicit guarantee of approximation error for any (non-random) pair of matrices $A$, $B$ in terms of only Frobenius norms $\|\bar{A}\|_F, \|\bar{B}\|_F$ and $\|\bar{A}^\top \bar{B}\|_F$, where $\bar{A},\bar{B}$ are versions of $A,B$ with zero-centered columns, respectively. For iid Gaussian matrices our quantizer achieves the lower bound and is, thus, asymptotically optimal. A practical low-complexity version of our quantizer achieves performance quite close to optimal. In addition, we derive rate-distortion function for matrix multiplication of iid Gaussian matrices, which exhibits an interesting phase-transition at $R\approx 0.906$ bit/entry, showing necessity of Johnson-Lindestrauss dimensionality reduction (sketching) in the low-rate regime.
academic

Cuantificación Óptima para Multiplicación de Matrices

Información Básica

  • ID del Artículo: 2410.13780
  • Título: Optimal Quantization for Matrix Multiplication
  • Autores: Or Ordentlich (Hebrew University of Jerusalem), Yury Polyanskiy (MIT)
  • Clasificación: cs.IT cs.AI cs.CL cs.LG math.IT
  • Fecha de Publicación: Octubre de 2024 (preimpresión en arXiv)
  • Enlace del Artículo: https://arxiv.org/abs/2410.13780

Resumen

Este artículo realiza un estudio profundo sobre el problema de cuantificación para multiplicación de matrices a gran escala. A diferencia de la cuantificación vectorial tradicional, el objetivo no es aproximar las matrices en sí, sino aproximar sus productos matriciales. Dadas dos matrices reales A y B, el codificador comprime cada matriz de forma independiente, describiendo cada entrada con R bits, y luego el decodificador utiliza estas representaciones comprimidas para estimar el producto matricial A⊤B. El artículo proporciona límites inferiores no asintóticos del error cuadrático medio aproximado para matrices con entradas gaussianas independientes e idénticamente distribuidas, construye cuantificadores universales basados en retículas anidadas, y descubre un fenómeno de transición de fase interesante en R ≈ 0.906 bits/entrada, indicando que en casos de tasa baja se requieren técnicas de reducción de dimensionalidad de Johnson-Lindenstrauss.

Antecedentes y Motivación de la Investigación

Definición del Problema

Con el auge de las redes neuronales profundas y los modelos de lenguaje grandes, la multiplicación de matrices se ha convertido en el principal cuello de botella computacional. El hardware computacional moderno suele estar limitado por el ancho de banda de memoria, no por la capacidad de cálculo. Por lo tanto, la compresión con pérdida de matrices para reducir la transferencia de memoria se convierte en un problema importante.

Necesidades Prácticas

Para modelos de lenguaje grandes, los autores estiman las tasas de cuantificación requeridas:

  • En la fase de generación, la CPU necesita una tasa de cuantificación de 1-3 bits/entrada para utilizar adecuadamente los recursos computacionales
  • En la fase de prefill, para LLMs pequeños ejecutándose en GPUs rápidas, se necesita una tasa de cuantificación de aproximadamente 11.7 bits/entrada

Limitaciones de Métodos Existentes

  1. Cuantificación Vectorial Clásica: La cuantificación independiente directa de matrices A y B, seguida del cálculo del producto de matrices cuantificadas, resulta en un error de O(n²)
  2. Métodos de Sketching: Aunque proporcionan estimaciones insesgadas, la varianza sigue siendo O(n²)
  3. Cuantificadores Deterministas: Existen límites inferiores de Ω(n²) para vectores en la esfera

Contribuciones Principales

  1. Límites Teóricos Inferiores: Proporciona límites inferiores no asintóticos para cuantificación de multiplicación de matrices con entradas gaussianas iid
  2. Cuantificador Universal: Construye un cuantificador universal basado en retículas anidadas con garantías de error explícitas para matrices arbitrarias
  3. Optimalidad Asintótica: Demuestra que el cuantificador propuesto alcanza el límite inferior para matrices gaussianas iid, siendo por lo tanto asintóticamente óptimo
  4. Fenómeno de Transición de Fase: Descubre una transición de fase en R ≈ 0.906 bits/entrada, revelando la necesidad de reducción de dimensionalidad en tasas bajas
  5. Algoritmo Práctico: Proporciona implementaciones de baja complejidad con rendimiento cercano al óptimo

Explicación Detallada de Métodos

Definición de la Tarea

Dadas matrices A ∈ R^(n×a) y B ∈ R^(n×b), el objetivo es diseñar codificadores f₁: R^(n×a) → 2^(naR) y f₂: R^(n×b) → 2^(nbR) así como un decodificador g, tales que:

E∥A⊤B−g(f1(A),f2(B))∥F2E\|A^⊤B - g(f_1(A), f_2(B))\|_F^2

sea minimizado, donde cada entrada de matriz se describe con R bits.

Función de Tasa-Distorsión Central Γ(R)

El artículo define la función de tasa-distorsión clave:

Γ(R)={1−(1−(2⋅2−2R∗−2−4R∗))RR∗R<R∗2⋅2−2R−2−4RR≥R∗\Gamma(R) = \begin{cases} 1 - \left(1 - (2 \cdot 2^{-2R^*} - 2^{-4R^*})\right) \frac{R}{R^*} & R < R^* \\ 2 \cdot 2^{-2R} - 2^{-4R} & R \geq R^* \end{cases}

donde R* ≈ 0.906 es la solución de la ecuación de punto fijo R = ½log₂(1 + 4R ln 2).

Esquema de Cuantificación con Retículas Anidadas

1. Cuantificación de Producto Interno (Bloque Constructivo Fundamental)

Para vectores U, V correlacionados con ρ en la esfera, se utiliza la retícula anidada Λc ⊂ Λf:

Proceso de Codificación:

  • Se añaden vectores de perturbación independientes Z₁, Z₂ a U y V respectivamente
  • Se cuantifica a la retícula fina Λf
  • Se genera la representación de coset en la retícula gruesa Λc

Proceso de Decodificación:

  • Se recupera el punto cuantificado del coset
  • Se elimina la perturbación
  • Se calcula la estimación del producto interno

2. Cuantificación de Multiplicación de Matrices

Pasos de Preprocesamiento:

  1. Centrado en Cero: Se calcula Ā = A - (1/n)1·1^⊤A, B̄ = B - (1/n)1·1^⊤B
  2. Cuantificación de Normas: Se describe con alta precisión la media y norma de cada columna
  3. Rotación Aleatoria: Se aplica la misma matriz ortogonal S a Ā y B̄

Pasos de Cuantificación:

  • Se aplica el cuantificador de producto interno a cada columna rotada
  • Se utilizan parámetros de compartición temporal κ y parámetro de escalado MMSE α

Puntos de Innovación Técnica

  1. Técnica de Perturbación: Hace que el error de cuantificación sea independiente de la entrada, evitando el error O(n²) de cuantificadores deterministas
  2. Estructura de Retícula Anidada: Proporciona tasa finita manteniendo buen rendimiento de cuantificación
  3. Compartición Temporal: Logra optimalidad en tasas bajas mediante reducción de dimensionalidad
  4. Rotación Aleatoria: Convierte vectores arbitrarios a distribución uniforme en la esfera, facilitando el análisis

Configuración Experimental

Experimentos de Verificación Teórica

Generación de Datos:

  • Matrices A, B ∈ R^(n×n) con entradas iid N(0,1)
  • n = 3 × 2¹¹

Detalles de Implementación:

  • Retícula base: retícula D₃ (3 dimensiones)
  • Razón de anidamiento: q = 6
  • Tamaño de tabla de búsqueda: < 64KB (cabe en caché L1)
  • Tasa efectiva: ≈ 3.015 bits/símbolo

Métodos de Comparación

  • Cuantificador escalar de 3 bits (normalización ℓ∞)
  • Límite teórico inferior Γ(R)

Resultados Experimentales

Resultados Principales

Comparación de Rendimiento:

  • Método propuesto: 1/n³ ∥Â⊤B - A⊤B∥²F ≈ 0.0593
  • Cuantificación escalar de 3 bits: ≈ 0.1668 (aproximadamente 3 veces de diferencia)
  • Límite teórico inferior: Γ(3.015) = 0.0304

Hallazgos Clave:

  1. El esquema basado en retícula D₃ supera significativamente la cuantificación escalar
  2. El rendimiento está cerca del óptimo teórico (aproximadamente 2 veces de diferencia)
  3. La brecha de rendimiento se reducirá aún más con el crecimiento de n

Análisis de Complejidad

Complejidad de Codificación: O(n log n) (usando transformada rápida de Hadamard) Complejidad de Decodificación: O(1) (usando tabla de búsqueda) Sobrecarga de Almacenamiento: Cada cuantificador requiere O(log n) bits adicionales para describir factores de escala

Trabajo Relacionado

Álgebra Lineal Aleatoria

  • Multiplicación de Matrices Monte Carlo (MCMM): Muestreo aleatorio de filas para aproximación
  • Hashing Sensible a Localidad (LSH): Para sketching de baja dimensión de similitud de coseno
  • Limitaciones: El error relativo crece con ∥A∥²F∥B∥²F/∥A⊤B∥²F

Cuantificación de Redes Neuronales

  • Cuantificación Post-Entrenamiento: Métodos como OPTQ, GPTQ
  • Técnicas de Rotación: QuIP, QuaRot utilizan transformadas de Hadamard
  • Cuantificación con Retículas: QUIP# utiliza retícula E₈ para cuantificación de pesos

Teoría de la Información

  • Compresión Distribuida: Compresión para computación de funciones lineales
  • Diseño de Códigos: Códigos de Voronoi y códigos de retícula anidada

Conclusiones y Discusión

Conclusiones Principales

  1. Optimalidad: Para matrices gaussianas iid, el esquema propuesto alcanza el límite de teoría de la información
  2. Universalidad: Proporciona garantías de rendimiento explícitas para matrices arbitrarias
  3. Fenómeno de Transición de Fase: R* ≈ 0.906 bits/entrada es el umbral crítico
  4. Practicidad: La implementación de baja complejidad se aproxima al rendimiento teórico

Limitaciones

  1. Aleatoriedad Compartida: Requiere que el codificador y decodificador compartan una semilla aleatoria
  2. Condiciones de Matriz: Requiere que las entradas de matriz estén acotadas (M = n^(10^22000))
  3. Retículas de Alta Dimensión: La optimalidad teórica requiere retículas "buenas" de alta dimensión; en la práctica se utilizan productos de retículas de baja dimensión
  4. Esquema Determinista: No está claro si existe un esquema determinista óptimo que no requiera aleatoriedad

Direcciones Futuras

  1. Multiplicación de Múltiples Matrices: Extensión a productos de k>2 matrices
  2. Otras Métricas de Distancia: Medidas no-MSE como divergencia KL
  3. Cuantificadores Deterministas: Exploración de esquemas sin necesidad de aleatoriedad compartida
  4. Aplicación en Redes Profundas: Despliegue y optimización en LLMs reales

Evaluación Profunda

Fortalezas

  1. Rigor Teórico: Proporciona análisis completo de teoría de la información, incluyendo límites superiores e inferiores
  2. Valor Práctico: Resuelve el cuello de botella real en inferencia de LLMs
  3. Innovación Técnica: Combina ingeniosamente cuantificación con retículas, rotación aleatoria y compartición temporal
  4. Universalidad: No depende de suposiciones de distribución de matriz específicas

Insuficiencias

  1. Complejidad: El análisis teórico es bastante complejo; la implementación real requiere múltiples componentes técnicos
  2. Factores Constantes: Aunque asintóticamente óptimo, los factores constantes en muestras finitas pueden ser grandes
  3. Adaptación de Hardware: Requiere optimización de implementación para diferentes plataformas de hardware
  4. Escalabilidad: La extensión de 2 matrices a múltiples productos matriciales no es trivial

Impacto

Contribuciones Teóricas:

  • Establece los fundamentos de teoría de la información para cuantificación de multiplicación de matrices
  • Revela el fenómeno de transición de fase y la necesidad de reducción de dimensionalidad
  • Proporciona un punto de referencia teórico importante para el campo

Aplicaciones Prácticas:

  • Proporciona nueva orientación teórica para cuantificación de LLMs
  • El trabajo posterior NestQuant ha logrado rendimiento SOTA en LLMs reales
  • Proporciona base teórica para diseño de aceleradores de hardware

Escenarios Aplicables

  1. Inferencia de Modelos de Lenguaje Grande: Cuantificación conjunta de pesos y activaciones
  2. Computación en el Borde: Operaciones matriciales en entornos con memoria limitada
  3. Computación Distribuida: Multiplicación de matrices con comunicación limitada
  4. Computación Científica: Problemas de álgebra lineal numérica a gran escala

Referencias

El artículo cita 44 referencias relacionadas, cubriendo múltiples campos incluyendo teoría de la información, teoría de retículas, álgebra lineal aleatoria y cuantificación de redes neuronales. Las referencias particularmente destacables incluyen:

  • La monografía de Zamir sobre codificación con retículas proporciona la base teórica
  • El trabajo pionero de Erez y Zamir sobre retículas anidadas
  • Métodos recientes de cuantificación de LLMs como OPTQ, QuIP, etc.
  • Resultados de teoría de matrices aleatorias y geometría esférica

Evaluación General: Este es un artículo excelente con contribuciones importantes tanto en teoría como en práctica, proporcionando una base sólida de teoría de la información para el problema de cuantificación de multiplicación de matrices, y demostrando un algoritmo práctico cercano al óptimo. Este trabajo tiene una importancia significativa para entender y mejorar técnicas de cuantificación en sistemas de aprendizaje automático a gran escala.