Локас: Ваши модели — это принципиальные инициализаторы локально поддерживаемых параметрических воспоминаний

Аннотация

В данной работе мы стремимся объединить обучение во время тестирования с новым типом параметрической памяти, которая может гибко выгружаться из параметров модели или объединяться с ними. Мы представляем Locas — локально поддерживаемую параметрическую память, которая использует архитектуру FFN-блоков современных трансформеров, что позволяет гибко интегрировать её в параметры модели, обеспечивая при этом эффективное непрерывное обучение. Мы рассматриваем две основные вариации Locas: первая использует традиционную двухслойную MLP-архитектуру с более строгими теоретическими гарантиями; вторая повторяет структуру GLU-FFN передовых языковых моделей и может быть легко подключена к существующим моделям для ресурсоэффективного (по параметрам и вычислениям) непрерывного обучения. Ключевым моментом является демонстрация того, что правильная инициализация такой низкоранговой памяти в стиле бокового FFN — выполненная принципиальным образом путём повторного использования параметров модели, активаций и/или градиентов — критически важна для быстрой сходимости, улучшенной обобщающей способности и предотвращения катастрофического забывания. Мы проверяем предложенный механизм памяти на задачах языкового моделирования целых книг PG-19 и ответов на вопросы в диалогах с длинным контекстом LoCoMo. Всего с 0.02% дополнительных параметров в минимальной конфигурации Locas-GLU способна сохранять информацию из предыдущего контекста, поддерживая при этом значительно меньшее окно контекста. Кроме того, мы оцениваем потерю общих способностей модели после запоминания всей книги с помощью Locas через сравнительное тестирование MMLU. Результаты демонстрируют перспективную способность Locas преобразовывать прошлый контекст в параметрическое знание с минимальным катастрофическим забыванием существующих внутренних знаний модели.

English

In this paper, we aim to bridge test-time-training with a new type of parametric memory that can be flexibly offloaded from or merged into model parameters. We present Locas, a Locally-Supported parametric memory that shares the design of FFN blocks in modern transformers, allowing it to be flexibly permanentized into the model parameters while supporting efficient continual learning. We discuss two major variants of Locas: one with a conventional two-layer MLP design that has a clearer theoretical guarantee; the other one shares the same GLU-FFN structure with SOTA LLMs, and can be easily attached to existing models for both parameter-efficient and computation-efficient continual learning. Crucially, we show that proper initialization of such low-rank sideway-FFN-style memories -- performed in a principled way by reusing model parameters, activations and/or gradients -- is essential for fast convergence, improved generalization, and catastrophic forgetting prevention. We validate the proposed memory mechanism on the PG-19 whole-book language modeling and LoCoMo long-context dialogue question answering tasks. With only 0.02\% additional parameters in the lowest case, Locas-GLU is capable of storing the information from past context while maintaining a much smaller context window. In addition, we also test the model's general capability loss after memorizing the whole book with Locas, through comparative MMLU evaluation. Results show the promising ability of Locas to permanentize past context into parametric knowledge with minimized catastrophic forgetting of the model's existing internal knowledge.