Усиленные быстрые веса с предсказанием следующей последовательности

Reinforced Fast Weights with Next-Sequence Prediction

February 18, 2026

Авторы: Hee Seung Hwang, Xindi Wu, Sanghyuk Chun, Olga Russakovsky

cs.AI

Аннотация

Архитектуры с быстрыми весами предлагают перспективную альтернативу трансформерам на основе внимания для моделирования длинных контекстов, обеспечивая постоянные затраты памяти независимо от длины контекста. Однако их потенциал ограничен парадигмой обучения предсказанию следующего токена. NTP оптимизирует предсказание одиночных токенов и игнорирует семантическую связность между несколькими токенами, следующими за префиксом. Как следствие, модели с быстрыми весами, которые динамически обновляют свои параметры для хранения контекстной информации, усваивают субоптимальные представления, не способные уловить дальнодействующие зависимости. Мы представляем REFINE — фреймворк обучения с подкреплением, который тренирует модели с быстрыми весами на основе целевой функции предсказания следующей последовательности. REFINE выбирает информативные позиции токенов на основе энтропии предсказания, генерирует rollout'ы из нескольких токенов, назначает самоконтролируемые вознаграждения на уровне последовательности и оптимизирует модель с помощью групповой относительной оптимизации политики. REFINE применим на всех этапах жизненного цикла предобученных языковых моделей: в процессе дообучения, после обучения и во время обучения на этапе тестирования. Наши эксперименты на LaCT-760M и DeltaNet-1.3B демонстрируют, что REFINE стабильно превосходит контролируемую тонкую настройку с NTP в задачах поиска иголки в стоге сена, вопросно-ответных системах с длинным контекстом и разнообразных задачах из LongBench. REFINE предоставляет эффективный и универсальный фреймворк для улучшения моделирования длинных контекстов в архитектурах с быстрыми весами.

English

Fast weight architectures offer a promising alternative to attention-based transformers for long-context modeling by maintaining constant memory overhead regardless of context length. However, their potential is limited by the next-token prediction (NTP) training paradigm. NTP optimizes single-token predictions and ignores semantic coherence across multiple tokens following a prefix. Consequently, fast weight models, which dynamically update their parameters to store contextual information, learn suboptimal representations that fail to capture long-range dependencies. We introduce REFINE (Reinforced Fast weIghts with Next sEquence prediction), a reinforcement learning framework that trains fast weight models under the next-sequence prediction (NSP) objective. REFINE selects informative token positions based on prediction entropy, generates multi-token rollouts, assigns self-supervised sequence-level rewards, and optimizes the model with group relative policy optimization (GRPO). REFINE is applicable throughout the training lifecycle of pre-trained language models: mid-training, post-training, and test-time training. Our experiments on LaCT-760M and DeltaNet-1.3B demonstrate that REFINE consistently outperforms supervised fine-tuning with NTP across needle-in-a-haystack retrieval, long-context question answering, and diverse tasks in LongBench. REFINE provides an effective and versatile framework for improving long-context modeling in fast weight architectures.