Por qué la eficiencia de E/S de memoria es importante para el rendimiento del modelo de IA

Tabla de enlaces

Resumen y 1 Introducción

2. Trabajo relacionado

3. Antecedentes

3.1. Notación y 3.2. Inferencia del modelo de lenguaje

3.3. Multi-Consulta, Multi-Encabezado y la Atención Multi-Consulta Generalizada

4. Atención bifurcada consciente del contexto y 4.1. Motivación

4.2. Formulación y 4.3. Complejidad de E/S de memoria

5. Experimentos

5.1. Comparación de las capacidades de atención multi-cabezal, multi-consulta y multi-grupo

5.2. Latencias de los modelos equivalentes a capacidades

5.3. Aplicaciones

6. Conclusión y referencias

A. Preguntas frecuentes

B. Trabajo relacionado

C. Configuración

D. Atención Multigrupal Familiar

E. Atención bifurcada consciente del contexto

F. Aplicaciones: Resultados adicionales

G. Compatibilidad con técnicas de decodificación especulativa y de decodificación rápida

B. Trabajo relacionado

B.1. Aplicaciones del muestreo por lotes en un único contexto

La reducción de latencia que logramos puede tener un profundo impacto en muchas aplicaciones. Algunas de estas aplicaciones incluyen:

• Generación de código: en el desarrollo de software, la generación de código asistida por IA puede beneficiarse enormemente de una latencia reducida, especialmente cuando se generan múltiples fragmentos de código o sugerencias para un contexto determinado. Esto puede generar una experiencia de usuario más ágil y eficiente para los desarrolladores que utilizan entornos de desarrollo integrados (IDE) o herramientas de finalización de código impulsados por IA (Nijkamp et al., 2023; 2022; Chen et al., 2021; Le et al., 2022; Fried et al., 2022; Li et al., 2022; Allal et al., 2023; Li et al., 2023; Ahmad et al., 2021).

• Traducción automática: en situaciones en las que se necesitan múltiples traducciones para una única entrada, como generar traducciones con distintos grados de formalidad o generar traducciones para diferentes dialectos, la atención bifurcada consciente del contexto puede proporcionar un cálculo más eficiente, lo que da como resultado servicios de traducción automática más rápidos y escalables (Costajussà et al., 2022; Farhad et al., 2021; Tran et al., 2021; Yee et al., 2019).

• Chatbots e IA conversacional: los agentes conversacionales a menudo necesitan generar múltiples respuestas para manejar diferentes interpretaciones de la entrada de un usuario o para brindar múltiples sugerencias. La latencia reducida que ofrece el método propuesto puede mejorar significativamente la capacidad de respuesta de los chatbots, lo que genera una conversación más natural y fluida con los usuarios (Google, 2023).

• Generación de contenido creativo: en aplicaciones como la creación de poesía, historias o publicidad, la capacidad de generar múltiples variaciones para un mensaje determinado es crucial. El método propuesto permite una generación más eficiente de contenido diverso, lo que lo hace más factible para aplicaciones en tiempo real o a gran escala (Lin y Riedl, 2021; Mirowski et al., 2023; Team, 2023; Yuan et al., 2022).

• Aumento de datos: en el contexto del aumento de datos para el aprendizaje automático, la generación de múltiples ejemplos alternativos para una entrada determinada puede ayudar a mejorar la solidez y la generalización del modelo. Con la latencia reducida que proporciona la atención bifurcada consciente del contexto, el proceso de generación de datos aumentados se puede acelerar, lo que permite un uso más eficiente de los recursos computacionales durante el entrenamiento.

• Evaluación general a gran escala: además de los casos de uso mencionados anteriormente, hay muchos casos de uso de nicho en los que se exploran LLM y otros modelos de generación abiertos en busca de toxicidad (Dathathri et al., 2019; Gehman et al., 2020; Nadeem et al., 2020), detección de código vulnerable en generaciones (Pearce et al., 2022), generación de edición de código que mejora el rendimiento (Madaan et al., 2023), traducciones de lenguajes de programación (Roziere et al., 2020) y muchos otros. En todos estos escenarios, se recopilan muchas generaciones por cada solicitud para una comprensión más profunda de los modelos; la atención bifurcada puede acelerar drásticamente el proceso de generación en tales casos.

En conclusión, el método de atención bifurcada consciente del contexto propuesto puede reducir significativamente el costo de E/S de memoria y mejorar la latencia en varias aplicaciones, lo que genera una mayor eficiencia y escalabilidad. Este método tiene el potencial de permitir nuevos casos de uso y mejorar la experiencia del usuario en numerosos sistemas impulsados por IA, haciéndolos más prácticos para la implementación en el mundo real.

B.2. El soporte de contextos extensos requiere atención eficiente en términos de E/S

A medida que los modelos de lenguaje se vuelven de propósito general y altamente capaces, la demanda de modelos de lenguaje para manejar secuencias de contexto más largas ha crecido significativamente. Recientemente, existe un enfoque continuo en modelos que pueden manejar secuencias de contexto aún más largas (Bulatov et al., 2023; OpenAI, 2023; Team, 2023). A partir de hoy, GPT-4 (OpenAI, 2023) admite una longitud de contexto de 32k tokens, y MPT-7B (Team, 2023) la extiende a 64k, mientras que Claude de Anthropic [3] admite una longitud de entrada de hasta 100k. Más recientemente, Bulatov et al propusieron una longitud de contexto de entrada de 1M de token para transformadores. Estos modelos amplían los límites de la comprensión del contexto y las capacidades de generación, lo que permite una comprensión del discurso más integral y respuestas informadas contextualmente.

Esta tendencia está impulsada por la necesidad de una comprensión integral del discurso en aplicaciones como la Generación Aumentada por Recuperación (RAG), así como muchos métodos complejos de incitación. Aplicaciones como RAG (Guu et al., 2020; Izacard et al., 2022; Menick et al., 2022; Zhen et al., 2022) recuperan pasajes o documentos extensos de corpus externos, lo que proporciona un contexto rico y fundamentado para generar respuestas. Además, modelos como Toolformer (Schick et al., 2023) y WebGPT (Nakano et al., 2021) aprovechan herramientas externas, como API y motores de búsqueda, para ampliar el contexto y mejorar la generación.

El contexto largo es desproporcionadamente costoso para los modelos de la familia de transformadores porque para la autoatención vainilla tanto la complejidad de la memoria como la del tiempo son cuadráticas a la longitud de la secuencia. Para manejar de manera efectiva secuencias de contexto más largas, es fundamental optimizar la E/S de memoria y reducir la sobrecarga computacional. Actualmente, los enfoques dominantes para abordar este desafío han sido hacer que el cálculo de la atención sea menos costoso. Beltagy et al. (2020) propusieron esparcir la autoatención utilizando varios patrones de atención. Wang et al. (2020) explora la aproximación de bajo rango de la autoatención. Además de las mejoras en los límites de cómputo, los avances en los mecanismos de atención eficientes en memoria y las técnicas para reducir la E/S de memoria seguirán impulsando el campo hacia adelante, facilitando el manejo de secuencias de contexto más largas en modelos de lenguaje. Se propone FlashAttention (Dao et al., 2022) para acelerar la autoatención y reducir la huella de memoria sin ninguna aproximación. Aprovecha el kernel fusionado para la multiplicación de matrices y la operación softmax, lo que reduce en gran medida la E/S de memoria durante el entrenamiento.

[3] https://www.anthropic.com/index/100k-context-windows