yifanzhang-pro/recurrent-looped-tranformer

해결하는 문제

Recurrent Looped Transformer (RLT)는 표준 Transformer가 가진 길이 일반화(학습 시 보지 못한 훨씬 긴 시퀀스에서 작업을 수행하는 능력)의 한계를 해결합니다. 특히 정확도를 유지하기 위해 시퀀스 길이에 따라 계산 경로가 커져야 하는 알고리즘 작업(패리티 및 상태 추적 등)을 대상으로 합니다.

작동 원리

RLT는 Transformer 아키텍처를 수정하여 토큰 간의 순환을 도입합니다:

• 순환 피드백 (Recurrent Feedback)： 이전 토큰의 최종 디코더 상태가 게이트 병합 (gated merge)을 통해 현재 토큰의 디코더 입력으로 피드백됩니다. 이는 시퀀스 길이에 따라 선형적으로 증가하는 계산 경로를 생성합니다.
• 하이브리드 메모리 (Hybrid Memory)： 글로벌 키-값 (KV) 메모리에는 병렬 인과 인코더를 사용하고, 로컬 메모리에는 슬라이딩 윈도우 어텐션 (SWA)이 있는 디코더를 사용합니다.
• 상태 연속성 (State Continuity)： 순환 상태는 프롬프트와 응답 사이의 경계를 매끄럽게 이어갑니다.
• 피드백 변형 (Feedback Variants)： 이 프로젝트는 다양한 피드백 간격 ($B$)을 탐구합니다: RLT-1은 매 토큰마다 업데이트, RLT-2는 $B$ 토큰 단위의 청크마다 업데이트, RLT-0은 피드백 없음.

대상 사용자

시퀀스 모델링, 길이 일반화, LLM의 알고리즘 추론을 연구하며 표준 Transformer의 고정 깊이 제한을 넘어서고자 하는 연구자 및 개발자.

주요 특징

• 강력한 길이 일반화： 학습 데이터보다 훨씬 긴 시퀀스(예: 패리티 작업의 경우 256비트)에서도 패리티 및 S5 스왑 작업에서 100% 정확도를 달성합니다.
• 순환 경로： 계산 경로는 $t$개의 토큰에 대해 $tL_D$개의 디코더 블록을 통과하며, 토큰당 비용을 일정하게 유지합니다.
• 유연한 깊이 할당： 인코더와 디코더 간의 레이어 분할이 다양한 알고리즘 작업의 성능에 어떤 영향을 미치는지 실험을 통해 보여줍니다.
• 효율성 트레이드오프： RLT-2(청크 피드백)는 일부 작업에서 높은 정확도를 유지하면서 RLT-1 대비 학습 속도를 향상(최대 2.27배)시킵니다.