메모리 효율적인 루프 트랜스포머: 루프 언어 모델에서 연산과 메모리를 분리하다

순환형 LLM 아키텍처는 중간 토큰을 생성하지 않고 임베딩 공간에서 다단계 계산을 가능하게 함으로써 추론 능력을 향상시키는 유망한 접근 방식으로 부상했습니다. Ouro와 같은 모델은 반복적으로 내부 표현을 업데이트하면서 표준 Key-Value (KV) 캐시를 반복 전반에 걸쳐 유지하며 추론함으로써, 메모리 사용량이 추론 깊이에 선형적으로 증가하는 특징이 있습니다. 결과적으로, 추론 반복 횟수를 늘리는 것은 막대한 메모리 사용으로 이어져 이러한 아키텍처의 실제 확장성을 제한할 수 있습니다. 본 논문에서는 추론 깊이를 메모리 소비와 분리하는 새로운 아키텍처인 Memory-Efficient Looped Transformer (MELT)를 제안합니다. MELT는 레이어 및 루프당 표준 KV 캐시를 사용하는 대신, 추론 루프 전반에 걸쳐 공유되는 단일 KV 캐시를 유지합니다. 이 캐시는 학습 가능한 게이팅 메커니즘을 통해 시간이 지남에 따라 업데이트됩니다. 이 아키텍처 하에서 안정적이고 효율적인 훈련을 가능하게 하기 위해, 우리는 LoopLM 시작 모델로부터 MELT로의 두 단계 절차인 보간 전이(interpolated transition)와 어텐션 정렬 증류(attention-aligned distillation)를 사용하여 MELT를 덩어리별(chunk-wise) 훈련하는 것을 제안합니다. 경험적으로, 우리는 사전 학습된 Ouro 매개변수로 미세 조정된 MELT 모델이 비슷한 크기의 표준 LLM보다 성능이 우수하면서도 메모리 사용량은 해당 모델과 비슷하고 Ouro보다 훨씬 작다는 것을 보여줍니다. 전반적으로, MELT는 LoopLM의 성능을 희생하지 않으면서 상수 메모리로 반복 추론을 달성합니다.

경량의 사후 훈련(post-training) 절차만을 사용하여.