Gefocuste Transformer: Contrastief Trainen voor Contextschaling

Focused Transformer: Contrastive Training for Context Scaling

July 6, 2023

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

cs.AI

Samenvatting

Grote taalmodelen hebben een uitzonderlijk vermogen om nieuwe informatie op een contextuele manier te integreren. Het volledige potentieel van een dergelijke aanpak wordt echter vaak beperkt door een tekortkoming in de effectieve contextlengte. Een oplossing voor dit probleem is het uitrusten van een aandachtlaag met toegang tot een extern geheugen, dat bestaat uit (sleutel, waarde)-paren. Naarmate het aantal documenten toeneemt, neemt echter de verhouding tussen relevante en irrelevante sleutels af, waardoor het model zich meer op de irrelevante sleutels richt. Wij identificeren een belangrijk probleem, het zogenaamde afleidingsprobleem, waarbij sleutels die aan verschillende semantische waarden zijn gekoppeld, kunnen overlappen, waardoor ze moeilijk te onderscheiden zijn. Om dit probleem aan te pakken, introduceren we de Focused Transformer (FoT), een techniek die een trainingsproces gebruikt dat is geïnspireerd op contrastief leren. Deze nieuwe aanpak verbetert de structuur van de (sleutel, waarde)-ruimte, waardoor de contextlengte kan worden uitgebreid. Onze methode maakt het mogelijk om bestaande, grootschalige modellen te fine-tunen om hun effectieve context te verlengen. Dit wordt gedemonstreerd door onze fine-tuning van 3B en 7B OpenLLaMA checkpoints. De resulterende modellen, die wij LongLLaMA noemen, tonen vooruitgang in taken die een lange context vereisen. We laten verder zien dat onze LongLLaMA-modellen een contextlengte van 256 k vakkundig beheren voor passkey-retrieval.

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.