다음 시퀀스 예측을 통한 강화된 고속 가중치

초록

빠른 가중치 아키텍처는 컨텍스트 길이에 관계없이 일정한 메모리 오버헤드를 유지함으로써 장문 컨텍스트 모델링에 있어 주의 기반 트랜스포머에 대한 유망한 대안을 제공합니다. 그러나 그 잠재력은 다음 토큰 예측(NTP) 훈련 패러다임에 의해 제한됩니다. NTP는 단일 토큰 예측을 최적화할 뿐 접두어 이후 여러 토큰에 걸친 의미적 일관성을 무시합니다. 그 결과, 컨텍스트 정보를 저장하기 위해 매개변수를 동적으로 업데이트하는 빠른 가중치 모델은 장거리 의존성을 포착하지 못하는 최적이 아닌 표현을 학습하게 됩니다. 본 연구에서는 강화 학습 프레임워크인 REFINE(Reinforced Fast weIghts with Next sEquence prediction)을 소개합니다. REFINE은 다음 시퀀스 예측(NSP) 목표 하에서 빠른 가중치 모델을 훈련시킵니다. REFINE은 예측 엔트로피를 기반으로 정보성이 높은 토큰 위치를 선택하고, 다중 토큰 롤아웃을 생성하며, 자기 지도 방식의 시퀀스 수준 보상을 할당하고, 그룹 상대 정책 최적화(GRPO)를 통해 모델을 최적화합니다. REFINE은 사전 훈련된 언어 모델의 전체 훈련 생애주기(중간 훈련, 사후 훈련, 테스트 시간 훈련)에 적용 가능합니다. LaCT-760M 및 DeltaNet-1.3B에 대한 실험 결과, REFINE은 바늘 더미 검색, 장문 컨텍스트 질의응답 및 LongBench의 다양한 작업에서 NTP를 사용한 지도 미세 조정을 지속적으로 능가하는 것으로 나타났습니다. 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.

다음 시퀀스 예측을 통한 강화된 고속 가중치

Reinforced Fast Weights with Next-Sequence Prediction

초록

Support