YOCO++: Enhancing YOCO with KV Residual Connections for Efficient LLM Inference

1. El Problema

La inferencia de Modelos de Lenguaje Grandes (LLM) basados en la arquitectura Transformer enfrenta un cuello de botella significativo en el consumo de memoria debido al caché de claves y valores (KV cache). Para mitigar esto, se han desarrollado métodos de compresión de KV cache entre capas (cross-layer), como YOCO (You Only Cache Once), que comparten los KV de una capa intermedia con las capas superiores, reduciendo el consumo de memoria en un 50%.

Sin embargo, estos métodos sufren de una degradación de rendimiento no despreciable. La causa principal es que la compartición directa de KV de una capa a otra debilita la capacidad del modelo, ya que las capas superiores pierden la capacidad de generar sus propias representaciones KV específicas, lo que limita el flujo de información y la expresividad del modelo.

2. Metodología: YOCO++

Los autores proponen YOCO++, una mejora sobre YOCO que busca recuperar la capacidad del modelo sin sacrificar la eficiencia. La metodología se basa en los siguientes pilares:

• Conexiones Residuales entre KV: En lugar de simplemente compartir los KV de la capa media con las capas superiores, YOCO++ introduce conexiones residuales ponderadas en las capas inferiores (la mitad inferior del modelo).
• Mecanismo de Combinación: Para cada capa $i$ en la mitad inferior (donde $i \in \{2, ..., L/2\}$), se calcula una combinación ponderada de los KV de la capa actual ($K_i, V_i$) y los KV de la capa inferior (base) ($K_1, V_1$):
  $$K'_i = \alpha_{i,1}^K K_1 + \alpha_{i,2}^K K_i$$
  $$V'_i = \alpha_{i,1}^V V_1 + \alpha_{i,2}^V V_i$$
  Donde $\alpha$ son pesos residuales aprendibles.
• Factor de Escalado ($\lambda$): Para facilitar el aprendizaje de estos pesos residuales y evitar que el modelo se quede atascado en el estado inicial (donde los pesos de la capa base son 0), se introduce un factor de escala fijo $\lambda > 1$. Esto permite que los pesos residuales se ajusten más fácilmente durante el entrenamiento.
• Optimización de I/O: A diferencia de otros métodos que combinan KV en las capas superiores (lo que aumentaría la sobrecarga de E/S en caché), YOCO++ realiza esta combinación en las capas inferiores y cacha los KV combinados ($K', V'$). Esto asegura que las capas superiores solo lean una única versión combinada, manteniendo la eficiencia de E/S idéntica a YOCO.

3. Contribuciones Clave

1. Arquitectura Híbrida Eficiente: Se introduce un mecanismo que combina la eficiencia de la compresión entre capas con la capacidad de recuperación de información mediante conexiones residuales, mejorando la capacidad del modelo sin aumentar la complejidad computacional significativa.
2. Mantenimiento de la Eficiencia: A diferencia de métodos anteriores como FusedKV, que introducen sobrecarga de E/S en caché, YOCO++ mantiene la misma eficiencia de inferencia que YOCO (misma latencia de prellenado y mismo throughput de decodificación).
3. Innovación en Pesos Residuales: La propuesta de utilizar una combinación de la capa base con la capa actual, junto con un factor de escalado para la inicialización, demuestra ser una estrategia efectiva para aprender representaciones ricas sin colapsar el rendimiento.
4. Código Abierto: Los autores han liberado su implementación para fomentar la reproducibilidad y el avance en el campo.

4. Resultados Experimentales

Los experimentos se realizaron utilizando modelos basados en TinyLlama (1.1B parámetros, 22 capas) entrenados desde cero en el conjunto de datos SlimPajama.

• Rendimiento (Precisión):
  • YOCO++ logró el estado del arte (SOTA) entre los métodos de compresión de KV cache entre capas con una tasa de compresión del 50%.
  • Superó al Transformer estándar en tareas de razonamiento de sentido común (HellaSwag, OBQA, WinoGrande, ARC, PIQA), logrando una precisión promedio de 48.99% frente a 48.37% del Transformer estándar.
  • Superó a otros métodos de compresión como FusedKV, FusedKV-Lite y YOCO base.
• Eficiencia:
  • Latencia de Prellenado (TTFT): Reducida en un 50% respecto al Transformer estándar (igual que YOCO), ya que permite salir temprano antes de calcular las capas superiores.
  • Throughput de Decodificación: Consistentemente superior al Transformer estándar debido a la menor memoria consumida, permitiendo tamaños de lote (batch sizes) más grandes.
  • Overhead Computacional: El costo adicional de combinar los KV es insignificante comparado con el mecanismo de atención.
• Estudios de Ablación:
  • Sin el factor de escalado ($\lambda$), el rendimiento cae a niveles similares a YOCO, confirmando que los pesos no se aprenden bien sin él.
  • Eliminar las conexiones residuales de las claves (keys) también degrada el rendimiento, demostrando que es necesario combinar tanto claves como valores.

5. Significado e Impacto

El trabajo de YOCO++ es significativo porque resuelve la dicotomía tradicional entre eficiencia y capacidad en la inferencia de LLMs.

• Superación del Paradigma de Compresión: Demuestra que es posible reducir el consumo de memoria del KV cache en un 50% sin sacrificar la precisión del modelo; de hecho, en este caso, se supera al modelo Transformer completo.
• Viabilidad para Despliegue: Al mantener la misma eficiencia de E/S y computación que YOCO, YOCO++ es una solución práctica y escalable para el despliegue de LLMs en entornos con recursos limitados (como dispositivos de borde o servidores con memoria restringida).
• Nueva Dirección de Investigación: Introduce la idea de "conexiones residuales entre KV" como una vía prometedora para mejorar la arquitectura de Transformers, sugiriendo que la información de las capas inferiores puede enriquecer dinámicamente a las capas superiores de manera más efectiva que la compartición estática.

En resumen, YOCO++ representa un avance crucial hacia LLMs más eficientes y potentes, logrando un equilibrio óptimo entre la reducción de memoria y la calidad de la generación de texto.