PACE comprime a 4 bits el estado de los LLM lineales y ahorra un tercio de memoria

Un equipo de investigadores ha propuesto PACE, un cuantizador postentrenamiento que almacena el estado recurrente de los modelos de atención lineal en cuatro bits y no requiere reentrenamiento. En cuatro modelos híbridos de entre 4B y 35B parámetros, la variante con cabezas mixtas de 4 y 8 bits es comparable a DAMP de 8 bits en los tres modelos más grandes, y ocupa cerca de dos tercios de su memoria. La versión en la que todas las cabezas usan cuatro bits, con el mismo volumen de memoria, obtiene mejores resultados que Q-Mamba.

Por qué el estado se convierte en el cuello de botella

En los modelos con atención lineal, el contexto no vive en una KV-cache que crece, sino en un estado recurrente de tamaño fijo. A primera vista, resulta cómodo. Pero el sistema de inferencia tiene que mantener una copia independiente de ese estado para cada secuencia que decodifica de forma simultánea. Cuantas más copias caben en la memoria del acelerador, más solicitudes se atienden en paralelo y, por tanto, más barato sale cada token.

De ahí el interés por cuantizar el propio estado. Los métodos existentes lo redondean después de cada token generado. El error de redondeo se acumula así paso a paso y, con cuatro bits, la precisión cae de forma notable. Los trabajos previos llegaban a la misma conclusión: o se sube la precisión en bits, o se reentrena el modelo para la cuantización. PACE demuestra que existe un tercer camino.

Dos observaciones en las que se apoya PACE

La primera se refiere a la frecuencia del redondeo. PACE redondea el estado solo una vez cada varios tokens y mantiene los tokens intermedios en un pequeño búfer. Los kernels por chunks ya funcionan de forma parecida en la fase de prefill: actualizan el estado una vez por chunk y no en cada token. Menos redondeos significan menos error acumulado.

La segunda observación atañe a la estructura del propio estado en las capas gated-delta. Los valores grandes no se reparten de manera uniforme, sino que se concentran en columnas. La razón es que algunos canales de la actualización que se escribe en cada token siguen siendo grandes durante muchos tokens. Por eso PACE usa una escala de cuantización común para toda la columna. Una sola entrada grande deja de perjudicar la precisión del resto de valores, que de otro modo tendrían que compartir escala con ella.

Un tercer elemento es la precisión mixta. Algunas cabezas retienen el error de redondeo durante más tiempo que otras. PACE mantiene precisamente esas en ocho bits y pasa el resto a cuatro.

• Método: postentrenamiento, sin reentrenar el modelo
• Formato objetivo: cuatro bits por elemento del estado y ocho para las cabezas de «error más persistente»
• Redondeo: una vez cada varios tokens; los tokens intermedios se guardan en un búfer
• Escala de cuantización: común para cada columna del estado
• Validación: cuatro modelos híbridos de entre 4B y 35B parámetros

Qué muestran las mediciones

El PACE mixto con cabezas de 4 y 8 bits quedó al nivel de DAMP de 8 bits en tres de los cuatro modelos, los más grandes. La memoria necesaria fue de aproximadamente dos tercios de la de DAMP. La variante de cuatro bits puros supera a Q-Mamba en la comparación con el mismo volumen de memoria. Por separado, los autores señalan que redondear cada varios tokens conserva la aceleración de la decodificación que aporta el búfer: el ahorro de memoria no se paga con una pérdida de velocidad.

Este último punto importa en la práctica. Una cuantización del estado que añade operaciones de redondeo en cada paso puede consumir la ganancia de compactar el batch. Aquí el redondeo sale del bucle crítico, lo que reduce ese riesgo.

Dónde resulta útil en la infraestructura

Las arquitecturas híbridas, en las que las capas lineales se alternan con la atención convencional, aparecen cada vez con más frecuencia en modelos abiertos y comerciales. Para quienes los sirven, PACE ofrece una palanca clara: reducir en aproximadamente un tercio la memoria del estado respecto a un esquema de 8 bits, sin reentrenar ni ajustar el modelo. El espacio liberado puede destinarse a un batch mayor o a una cola más larga de sesiones simultáneas.

También hay limitaciones. El resultado se ha confirmado en cuatro modelos híbridos y, en el más pequeño de ellos, el PACE mixto no alcanzó el nivel de DAMP. Los equipos que planeen adoptarlo deberían comprobar cómo se comporta el método con su modelo y su carga antes de reorganizar la producción en torno a un estado de cuatro bits. Aun así, la idea en sí, redondear con menos frecuencia y ajustar la escala a la estructura real de los datos, parece trasladable y barata de implementar.