Longhorn: 상태 공간 모델은 분할 상환 방식의 온라인 학습자입니다

Longhorn: State Space Models are Amortized Online Learners

July 19, 2024

저자: Bo Liu, Rui Wang, Lemeng Wu, Yihao Feng, Peter Stone, Qiang Liu

cs.AI

초록

대형 언어 모델(LLMs)과 같은 현대 AI 방법론의 가장 근본적인 능력은 긴 토큰 시퀀스에서 다음 토큰을 예측하는 능력, 즉 "시퀀스 모델링"입니다. 트랜스포머(Transformers) 모델이 현재 시퀀스 모델링의 주류 접근법이지만, 시퀀스 길이에 대한 2차 계산 비용은 중요한 단점으로 지적됩니다. 상태-공간 모델(SSMs)은 선형 디코딩 효율성과 훈련 중 높은 병렬화 가능성으로 인해 유망한 대안으로 주목받고 있습니다. 그러나 기존 SSM들은 종종 임의적으로 보이는 선형 회귀 설계에 의존하는 경향이 있습니다. 본 연구에서는 SSM 설계를 온라인 학습의 관점에서 탐구하며, SSM을 특정 온라인 학습 문제를 위한 메타 모듈로 개념화합니다. 이 접근법은 SSM 설계를 정확한 온라인 학습 목표 설정과 연결하며, 이러한 목표를 최적화함으로써 상태 전이 규칙을 도출합니다. 이러한 통찰을 바탕으로, 우리는 온라인 회귀 목표를 최적화하기 위한 암시적 업데이트(implicit update)를 기반으로 한 새로운 심층 SSM 아키텍처를 제안합니다. 실험 결과, 우리의 모델은 표준 시퀀스 모델링 벤치마크와 언어 모델링 작업에서 Mamba 모델을 포함한 최첨단 SSM들을 능가하는 성능을 보여줍니다.

English

The most fundamental capability of modern AI methods such as Large Language Models (LLMs) is the ability to predict the next token in a long sequence of tokens, known as ``sequence modeling." Although the Transformers model is the current dominant approach to sequence modeling, its quadratic computational cost with respect to sequence length is a significant drawback. State-space models (SSMs) offer a promising alternative due to their linear decoding efficiency and high parallelizability during training. However, existing SSMs often rely on seemingly ad hoc linear recurrence designs. In this work, we explore SSM design through the lens of online learning, conceptualizing SSMs as meta-modules for specific online learning problems. This approach links SSM design to formulating precise online learning objectives, with state transition rules derived from optimizing these objectives. Based on this insight, we introduce a novel deep SSM architecture based on the implicit update for optimizing an online regression objective. Our experimental results show that our models outperform state-of-the-art SSMs, including the Mamba model, on standard sequence modeling benchmarks and language modeling tasks.