Resumen Técnico: Regularizar o Localizar: Cuándo la Geometría del K/V-Cache en el Tiempo de Entrenamiento Beneficia la Cuantización

Declaración del Problema

Los modelos de lenguaje Transformer entrenados con entropía cruzada (CE) suelen desarrollar una geometría de estados ocultos anisotrópica, donde las representaciones se concentran a lo largo de unas pocas direcciones compartidas. Los modelos de Arquitectura Predictiva de Embebido Conjunto (JEPA) han introducido objetivos como SIGReg (una prueba de normalidad de Epps–Pulley esquematizada) para empujar los embebidos hacia la isotropía ($N(0, I)$), pero los intentos previos de aplicar SIGReg a Grandes Modelos de Lenguaje (LLMs) han arrojado resultados nulos. Estos fallos ocurrieron cuando SIGReg se aplicó como un auxiliar a flujos no estándar o únicamente a los estados de la última capa durante el ajuste fino (fine-tuning).

Persiste una brevedad crítica: ¿Resetea SIGReg la geometría del preentrenamiento autorregresivo (AR) estándar, y más importante aún, beneficia esta reconfiguración al K/V-cache durante la cuantización post-hoc? El trabajo existente se centra en la compresión post-hoc o en cambios arquitectónicos, dejando sin explorar en gran medida si la regularización distributiva en el tiempo de entrenamiento de la geometría estándar del K/V-cache puede mejorar la robustez de la cuantización.

Metodología

Los autores emplean un diseño experimental de tres etapas utilizando modelos de estilo Llama de 110 millones de parámetros preentrenados con 10 mil millones de tokens de FineWeb.

1. Fase 1: Respuesta a la Dosis y Caracterización del Estado Oculto

• Objetivo: Determinar si añadir SIGReg a la pérdida CE durante el preentrenamiento reduce la anisotropía del estado oculto.
• Configuración: Los autores realizan un barrido del peso de regularización $\lambda$ ($0, 0.001, 0.01, 0.1, 1.0$).
• Métrica: Anisotropía de coseno por pares y perplejidad.
• Hallazgo: Se identifica $\lambda=0.01$ como el punto de operación óptimo, reduciendo significativamente la anisotropía con un coste de perplejidad mínimo.

2. Fase 2: Sondas de Puente (Adaptación vs. Retroajuste)

Para determinar si el cambio de geometría se propaga al K/V-cache y cómo lograrlo, los autores prueban tres sondas específicas partiendo de un checkpoint de $\lambda=0.01$:

• A1 (Adaptación): Descongelar el tronco (trunk) y continuar el entrenamiento aplicando SIGReg directamente tanto a los estados ocultos ($h$) como a las proyecciones K/V ($K_\ell, V_\ell$) en todas las capas.
• A2 (Predicción Congelada): Entrenar predictores por capa sobre un tronco congelado para ver si la geometría K/V es recuperable únicamente a partir de los estados ocultos.
• A3 (Decodificación Congelada): Entrenar decodificadores para reconstruir K/V a partir del estado oculto final del tronco congelado.
• Control: Una "continuación emparejada" donde el tronco se continúa durante el mismo presupuesto de tokens pero sin el término de regularización K/V ($\lambda_{KV}=0$).

3. Fase 3: Evaluación de la Cuantización

Los autores evalúan los modelos resultantes bajo cuantización falsa post-hoc (cuantizar $\to$ descuantizar) del K/V-cache.

• Esquemas Probados: Cuantización simétrica sin grupos (por token vs. por canal), rejillas asimétricas (con puntos de cero) y arreglos mixtos (estilo KIVI: por canal K, por token V).
• Métrica: $\Delta$NLL (Log-Verosimilitud Negativa) respecto a su propia referencia de precisión completa.

Resultados Clave

1. Geometría del Estado Oculto

En $\lambda=0.01$, SIGReg reduce la anisotropía de coseno por pares en un 38% a través de tres semillas emparejadas. La perplejidad aumenta menos del 0.35%, y las capacidades de zero-shot (LAMBADA, HellaSwag, etc.) permanecen limpias sin signos de degradación consistente.

2. Geometría del K/V-cache y el Fallo del "Retroajuste"

• No hay Propagación Automática: La reducción de la anisotropía en los estados ocultos no se propaga automáticamente al K/V-cache. Las sondas de tronco congelado (A2 y A3) fallan en recuperar la geometría, indicando que la anisotropía K/V no es meramente una función del estado oculto final.
• Se Requiere Regularización Directa: Solo A1 (entrenamiento continuado con regularización K/V directa) logra reducir con éxito la anisotropía del cache, alcanzando una reducción del 94% en la anisotropía media del cache.
• Aislamiento del Mecanismo: La continuación emparejada sin el término K/V ($\lambda_{KV}=0$) no logra reproducir este efecto, aislando el término K/V como la causa necesaria.

3. Recompensa de la Cuantización

Los beneficios de la geometría regularizada dependen altamente del esquema de cuantización:

• Escalas Gruesas (Sin Grupos): Bajo cuantización simétrica sin grupos y sin transformaciones, el modelo adaptado por A1 es el único que prefiere el escalado por canal sobre el escalado por token. Bajo cuantificación simétrica de 3 bits por canal, el modelo base sufre un daño de $\Delta$NLL de 4.3–7.9× más que el modelo A1.
• Escalas Finas (Estilo KIVI): Bajo una configuración completa simulada de estilo KIVI (arreglo mixto, puntos de cero y escalas agrupadas con $g=32$), la ventaja desaparece. Todos los modelos, incluido el base, alcanzan una paridad casi total en el rendimiento.
• Sobrecarga de Almacenamiento: Incluso al contabilizar la sobrecarga de metadatos de los puntos de cero agrupados (haciendo que el almacenamiento agrupado de 3 bits sea comparable al de 4 bits sin grupos), la ventaja del tiempo de entrenamiento desaparece bajo la configuración de escalado mixto y agrupación local por token.

Significancia y Reivindicaciones

El artículo afirma ser el primero en evaluar la regularización distributiva en el tiempo de entrenamiento de la geometría estándar del K/V-cache frente a la cuantización del cache post-hoc.

• Control Geométrico: Los autores demuestran que la geometría del K/V-cache puede controlarse explícitamente durante el entrenamiento mediante regularización directa, pero este efecto no surge únicamente de la regularización del estado oculto y no puede ser "retroajustado" mediante sondas de tronco congelado.
• Utilidad Condicional: La intervención en el tiempo de entrenamiento proporciona una recompensa tangible solo cuando los escaladores de cuantización son gruesos (abarcan extensiones de tokens largas). La ventaja es específica de la interacción entre la geometría regularizada y el esquema de cuantización; no se generaliza a las configuraciones modernas y finas (como el estilo KIVI con puntos de cero y agrupación local).
• Interpretación de Métricas: El artículo aclara que la métrica de anisotropía de coseno por pares mide la alineación de la dirección media (dominada por medias compartidas) en lugar del blanqueamiento de la covarianza. Los resultados muestran una disociación: mientras que la alineación de la media cae, las propiedades espectrales (como la participación del top-1) se comportan de forma distinta para las capas K y V, sugiriendo que el regularizador actúa principalmente sobre la media/direcciones errantes.

Limitaciones

Los autores delimitan explícitamente sus afirmaciones:

• Régimen: Los resultados se limitan a 110 millones de parámetros y al corpus FineWeb; modelos más grandes con diferentes disposiciones de cache (GQA/MQA) no han sido probados.
• Simulación: La cuantización es simulada (cuantificación falsa) sin mediciones de tiempo real o de memoria.
• Alcance de las Sondas: Los resultados negativos para las sondas congeladas se basan en capacidades y presupuestos específicos; un "punto medio" (por ejemplo, adaptación a escala LoRA) sigue sin probarse.
• No es la "Primera Isotropía": Los autores no pretenden ser los primeros en regularizar para la isotropía (citando CosReg, SimCTG, etc.), sino los primeros en apuntar específicamente a la geometría del K/V-cache para la robustez de la cuantización.