Los investigadores presentan JoLT, un método para comprimir la caché de clave-valor (KV) en la inferencia de transformadores tratándola como un tensor de tercer orden. El enfoque aplica una descomposición de Tucker parcial a los ejes de tokens y características, utilizando residuos rotados de bajo bit de Johnson-Lindenstrauss para restaurar la energía perdida durante la truncación.

  • Un único dual lagrangiano asigna rangos de Tucker y anchos de bits de residuos por grupo de capas bajo un presupuesto de un solo byte.
  • JoLT logra una compresión de 2-3x con perplexidad, precisión en GSM8K y recuperación RULER manteniéndose dentro del ruido estadístico de la línea base sin comprimir en Mistral-7B-v0.3 y LLaMA-2-13B.
  • A una compresión de 2x, el error relativo de Frobenius es 0.009 para las claves y 0.006 para los valores, un orden de magnitud menor que la SVD entre capas y la cuantización de 4 bits.
  • Una variante de SVD aleatorizada llamada FlashJoLT proporciona una aceleración del tiempo de compresión de 5-13x con calidad equivalente.

Este método aborda el costo dominante de memoria de la inferencia de transformadores reduciendo significativamente el tamaño de la caché KV sin degradar el rendimiento del modelo.