Focused Transformer: Контрастное обучение для масштабирования контекста

Focused Transformer: Contrastive Training for Context Scaling

July 6, 2023

Авторы: Szymon Tworkowski, Konrad Staniszewski, Mikołaj Pacek, Yuhuai Wu, Henryk Michalewski, Piotr Miłoś

cs.AI

Аннотация

Крупные языковые модели обладают исключительной способностью контекстуально усваивать новую информацию. Однако полный потенциал такого подхода часто ограничивается из-за недостаточной длины эффективного контекста. Одним из решений этой проблемы является предоставление слою внимания доступа к внешней памяти, состоящей из пар (ключ, значение). Тем не менее, по мере увеличения количества документов доля релевантных ключей по сравнению с нерелевантными уменьшается, что заставляет модель больше фокусироваться на нерелевантных ключах. Мы выявляем значительную проблему, названную проблемой отвлечения, при которой ключи, связанные с различными семантическими значениями, могут перекрываться, что затрудняет их различение. Для решения этой проблемы мы представляем Focused Transformer (FoT) — метод, использующий процесс обучения, вдохновленный контрастивным обучением. Этот новый подход улучшает структуру пространства (ключ, значение), позволяя расширить длину контекста. Наш метод позволяет тонко настраивать существующие крупномасштабные модели для увеличения их эффективного контекста. Это демонстрируется нашей тонкой настройкой контрольных точек OpenLLaMA на 3B и 7B параметров. Полученные модели, которые мы называем LongLLaMA, показывают улучшения в задачах, требующих длинного контекста. Мы также показываем, что наши модели LongLLaMA успешно справляются с длиной контекста в 256 тысяч токенов для задачи извлечения ключа доступа.

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.