Focused Transformer: Addestramento Contrastivo per il Ridimensionamento del Contesto

Focused Transformer: Contrastive Training for Context Scaling

July 6, 2023

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

cs.AI

Abstract

I modelli linguistici di grandi dimensioni possiedono una capacità eccezionale di incorporare nuove informazioni in modo contestuale. Tuttavia, il pieno potenziale di tale approccio è spesso limitato da una restrizione nella lunghezza effettiva del contesto. Una soluzione a questo problema consiste nel dotare uno strato di attenzione di accesso a una memoria esterna, composta da coppie (chiave, valore). Tuttavia, all'aumentare del numero di documenti, la proporzione di chiavi rilevanti rispetto a quelle irrilevanti diminuisce, portando il modello a concentrarsi maggiormente sulle chiavi irrilevanti. Identifichiamo una sfida significativa, denominata problema della distrazione, in cui le chiavi associate a diversi valori semantici potrebbero sovrapporsi, rendendole difficili da distinguere. Per affrontare questo problema, introduciamo il Focused Transformer (FoT), una tecnica che utilizza un processo di addestramento ispirato all'apprendimento contrastivo. Questo approccio innovativo migliora la struttura dello spazio (chiave, valore), consentendo un'estensione della lunghezza del contesto. Il nostro metodo permette di ottimizzare modelli preesistenti su larga scala per estendere il loro contesto effettivo. Questo è dimostrato dalla nostra ottimizzazione dei checkpoint OpenLLaMA da 3B e 7B. I modelli risultanti, che abbiamo chiamato LongLLaMA, mostrano progressi in compiti che richiedono un contesto lungo. Illustriamo inoltre che i nostri modelli LongLLaMA gestiscono abilmente una lunghezza del contesto di 256 k per il recupero di passkey.

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.