World Model on Million-Length Video And Language With Blockwise RingAttention

Enabling long-context understanding remains a key challenge in scaling existing sequence models -- a crucial component in developing generally intelligent models that can process and operate over long temporal horizons that potentially consist of millions of tokens. In this paper, we aim to address these challenges by providing a comprehensive exploration of the full development process for producing 1M context language models and video-language models, setting new benchmarks in language retrieval and new capabilities in long video understanding. We detail our long context data curation process, progressive context extension from 4K to 1M tokens, and present an efficient open-source implementation for scalable training on long sequences. Additionally, we open-source a family of 7B parameter models capable of processing long text documents and videos exceeding 1M tokens.

한 줄 요약

1M 토큰 길이의 언어 및 비디오-언어 모델을 구축하고, Blockwise RingAttention을 기반으로 확장성을 보장한 오픈소스 구현을 제시한다.

핵심 기여도

1M 토큰 길이의 언어 및 비디오-언어 모델(LWM)을 개발하여 언어 검색 및 장기 비디오 이해에서 새로운 벤치마크를 설정.
Blockwise RingAttention 기법을 활용해 4K에서 1M 토큰까지 점진적으로 확장하며, 계산 비용을 최소화한 훈련 전략 제시.
마스킹된 시퀀스 패킹과 손실 균형 조절을 통해 다중 모달(텍스트, 이미지, 비디오) 학습을 효과적으로 수행.
7B 파라미터의 LWM-Text, LWM-Chat, LWM, LWM-Video 모델을 오픈소스로 공개.

핵심 아이디어

기존 시퀀스 모델은 수천 토큰 이하의 짧은 문맥만 처리할 수 있어, 수백만 토큰 길이의 문서나 비디오를 처리하는 데 한계가 있었다. 본 연구는 Blockwise RingAttention이라는 새로운 어텐션 메커니즘을 도입하여, 1M 토큰 길이의 시퀀스를 효율적으로 처리할 수 있는 모델을 구축한다. 이 기법은 기존의 어텐션 확장 방식과 달리 근사치 없이 정확한 계산을 유지하면서도 메모리 사용량을 줄인다. 또한, 장기 텍스트와 비디오 데이터를 위한 커리어 프로세스를 구축하고, 텍스트 기반 모델이 생성한 합성 질문-답변 데이터를 활용해 대화 능력을 향상시킨다.

기술적 접근법

**모델 아키텍처**: Llama2 7B 기반의 오토리그레시브 트랜스포머를 사용.
**훈련 전략**: 2단계 훈련 방식을 채택.
**Stage I**: 4K에서 1M 토큰까지 점진적으로 확장.
**Stage II**: 텍스트-이미지, 텍스트-비디오, 장문 텍스트 데이터를 포함한 다중 모달 학습.
**데이터셋**: 공개 자료에서 추출한 장문 비디오와 책 데이터, 합성 QA 데이터를 사용.
**기타 기술**: 마스킹된 시퀀스 패킹, 손실 균형 조절, VQGAN 인코더를 통한 이미지/비디오 생성 손실 최적화.

주요 결과

**LOFT 데이터셋 평가**: 512K 토큰 길이에서 기존 언어 모델 베이스라인 대비 우수한 성능을 보임.
**장기 비디오 이해**: 1시간 이상의 비디오 처리 능력을 갖춘 LWM 모델 제시.
**장문 문서 검색**: 1M 토큰 길이의 문서에서 정확한 정보 검색 성능을 달성.
**모델 가속도**: Blockwise RingAttention을 통해 1.6× 이상의 훈련 가속 효과를 기록 (명시되지 않음).

의의 및 한계

본 연구는 수백만 토큰 길이의 시퀀스를 처리할 수 있는 모델을 구축함으로써, 장기 시퀀스 이해를 위한 기초를 마련했다. 특히, Blockwise RingAttention은 기존 어텐션 기법의 한계를 극복하며, 대규모 데이터셋에서의 확장성을 보여준다. 또한, 다중 모달 학습을 위한 마스킹 및 손실 균형 기법은 실제 적용 가능성에 기여한다.

하지만, 연구에는 몇 가지 한계가 있다. 첫째, 이미지 토크나이저는 단순한 방식을 사용하여, 시간적 중복성을 고려한 비디오 토크나이저 개선이 필요하다. 둘째, 모델 파라미터 수(7B)는 대형 언어 모델(100B 이상)에 비해 작아, 대규모 확장 시 다른 스케일링 특성이 나타날 수 있다.

실용적 활용

LWM 모델은 장문 문서 분석, 장기 비디오 요약, 대화형 챗봇 등에서 활용 가능하다. 특히, 1M 토큰 길이의 텍스트와 비디오를 처리할 수 있어, 법적 문서 검색, 교육 콘텐츠 요약, 영상 기반 QA 시스템 등에 적용할 수 있다. 오픈소스 구현을 통해 연구자 및 엔지니어가 장기 시퀀스 모델 개발에 쉽게 접근할 수 있다.