Transformador Enfocado: Entrenamiento Contrastivo para Escalado de Contexto

Focused Transformer: Contrastive Training for Context Scaling

July 6, 2023

Autores: Szymon Tworkowski, Konrad Staniszewski, Mikołaj Pacek, Yuhuai Wu, Henryk Michalewski, Piotr Miłoś

cs.AI

Resumen

Los modelos de lenguaje de gran escala tienen una capacidad excepcional para incorporar nueva información de manera contextual. Sin embargo, el potencial completo de este enfoque a menudo se ve limitado por una restricción en la longitud efectiva del contexto. Una solución a este problema es dotar a una capa de atención con acceso a una memoria externa, que consta de pares (clave, valor). No obstante, a medida que aumenta el número de documentos, la proporción de claves relevantes frente a las irrelevantes disminuye, lo que lleva al modelo a enfocarse más en las claves irrelevantes. Identificamos un desafío significativo, denominado el problema de distracción, donde las claves vinculadas a diferentes valores semánticos pueden solaparse, dificultando su distinción. Para abordar este problema, presentamos el Transformer Enfocado (FoT), una técnica que emplea un proceso de entrenamiento inspirado en el aprendizaje contrastivo. Este enfoque novedoso mejora la estructura del espacio (clave, valor), permitiendo una extensión de la longitud del contexto. Nuestro método permite ajustar modelos preexistentes de gran escala para ampliar su contexto efectivo. Esto se demuestra mediante nuestro ajuste fino de los puntos de control OpenLLaMA de 3B y 7B. Los modelos resultantes, que denominamos LongLLaMA, muestran avances en tareas que requieren un contexto largo. Además, ilustramos que nuestros modelos LongLLaMA gestionan hábilmente una longitud de contexto de 256k para la recuperación de claves de acceso.

English

Large language models have an exceptional capability to incorporate new information in a contextual manner. However, the full potential of such an approach is often restrained due to a limitation in the effective context length. One solution to this issue is to endow an attention layer with access to an external memory, which comprises of (key, value) pairs. Yet, as the number of documents increases, the proportion of relevant keys to irrelevant ones decreases, leading the model to focus more on the irrelevant keys. We identify a significant challenge, dubbed the distraction issue, where keys linked to different semantic values might overlap, making them hard to distinguish. To tackle this problem, we introduce the Focused Transformer (FoT), a technique that employs a training process inspired by contrastive learning. This novel approach enhances the structure of the (key, value) space, enabling an extension of the context length. Our method allows for fine-tuning pre-existing, large-scale models to lengthen their effective context. This is demonstrated by our fine-tuning of 3B and 7B OpenLLaMA checkpoints. The resulting models, which we name LongLLaMA, exhibit advancements in tasks requiring a long context. We further illustrate that our LongLLaMA models adeptly manage a 256 k context length for passkey retrieval.