Transformer Focado: Treinamento Contrastivo para Escalonamento 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

Resumo

Modelos de linguagem de grande escala possuem uma capacidade excepcional de incorporar novas informações de maneira contextual. No entanto, o potencial total dessa abordagem é frequentemente limitado devido a uma restrição no comprimento efetivo do contexto. Uma solução para esse problema é dotar uma camada de atenção com acesso a uma memória externa, que consiste em pares (chave, valor). Contudo, à medida que o número de documentos aumenta, a proporção de chaves relevantes em relação às irrelevantes diminui, fazendo com que o modelo se concentre mais nas chaves irrelevantes. Identificamos um desafio significativo, denominado problema de distração, no qual chaves associadas a diferentes valores semânticos podem se sobrepor, tornando-as difíceis de distinguir. Para enfrentar esse problema, introduzimos o Transformer Focado (FoT), uma técnica que emprega um processo de treinamento inspirado no aprendizado contrastivo. Essa nova abordagem melhora a estrutura do espaço (chave, valor), permitindo uma extensão do comprimento do contexto. Nosso método possibilita o ajuste fino de modelos pré-existentes e de grande escala para aumentar seu contexto efetivo. Isso é demonstrado pelo nosso ajuste fino dos checkpoints OpenLLaMA de 3B e 7B. Os modelos resultantes, que nomeamos LongLLaMA, exibem avanços em tarefas que exigem um contexto longo. Além disso, ilustramos que nossos modelos LongLLaMA gerenciam habilmente um comprimento de contexto de 256 k para recuperação de chaves de acesso.

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.