En palabras sencillas
Los modelos de lenguaje pueden tardar más en responder cuando tienen que revisar una conversación o un documento muy largo. Una forma de ahorrar trabajo es hacer que el modelo consulte solo algunos tokens anteriores, pero elegir únicamente los de mayor puntuación puede hacer que se pierda información útil.
Este artículo presenta Stratified Sparse Attention, que conserva los tokens con puntuaciones altas y toma muestras de otros tokens según estimaciones de su importancia. En una prueba de contexto largo, superó a un método práctico que toma muestras uniformes del resto de los tokens, aunque un método de selección basado en puntuaciones exactas obtuvo mejores resultados en una de las condiciones informadas.
La investigación
Cuando un modelo de lenguaje procesa un contexto largo, la atención dispersa puede reducir las consultas repetidas a una caché que va creciendo, pero elegir solo los tokens con mayor puntuación puede hacer que se pierda atención útil, y el muestreo uniforme puede desperdiciar selecciones. ¿Puede un método conservar más contexto útil y seguir atendiendo solo a una fracción de los tokens?
Stratified Sparse Attention combina la conservación de tokens con puntuaciones altas y el muestreo guiado por puntuaciones para captar más contexto útil con atención dispersa.
Los autores preguntan cómo puede la atención dispersa conservar información útil de un contexto largo sin leer todos los tokens almacenados en caché. Stratified Sparse Attention combina la conservación determinista de tokens con puntuaciones altas y el muestreo de tokens residuales guiado por puntuaciones proxy; después, corrige los pesos de los tokens muestreados según sus probabilidades de inclusión.
En RULER-HARD, SSA obtuvo una puntuación superior a la de una alternativa híbrida práctica que usa las mismas puntuaciones proxy y el mismo presupuesto de atención, pero muestrea los tokens residuales de forma uniforme.
La comparación respalda la interpretación de los autores de que el muestreo guiado por puntuaciones ayuda en estas condiciones; aun así, un método top-k con puntuaciones exactas obtuvo una puntuación superior en una de las condiciones informadas.
Estos resultados abarcan modelos, pruebas y longitudes de contexto seleccionados. Con Ministral-3-8B y un contexto de 128K, Oracle-top-k obtuvo una puntuación superior a SSA en RULER-HARD.
Además, la implementación práctica para GPU usa una fracción fija de tokens con puntuaciones altas, en lugar de la asignación adaptativa del algoritmo formal, y limita las probabilidades de inclusión de una manera que, según los autores, introduce un pequeño sesgo.
Análisis
Detalles
- Autores del paper original
- Sahil Joshi, Prithvi Dixit, Aryan Katiyar, Agniva Chowdhury, Joseph E. Gonzalez
- Enviado a arXiv
- Licencia
- CC BY 4.0