https://arxiv.org/abs/2407.04620

자기 주의는 긴 맥락에서 좋은 성능을 보이지만 2차 복잡도를 갖습니다. 기존 RNN 계층은 선형 복잡도를 갖지만 긴 맥락에서의 성능은 숨겨진 상태의 표현력에 의해 제한됩니다. 선형 복잡도와 표현력 있는 숨겨진 상태를 갖는 새로운 종류의 시퀀스 모델링 계층을 제안합니다. 핵심 아이디어는 숨겨진 상태를 그 자체로 머신 러닝 모델로 만들고 업데이트 규칙을 자체 감독 학습의 단계로 만드는 것입니다. 숨겨진 상태는 테스트 시퀀스에서도 학습을 통해 업데이트되므로 계층을 테스트 시간 학습(TTT) 계층이라고 합니다. 숨겨진 상태가 각각 선형 모델과 2계층 MLP인 TTT-Linear와 TTT-MLP의 두 인스턴스를 고려합니다. 강력한 Transformer와 최신 RNN인 Mamba와 비교하여 125M에서 1.3B 매개변수의 규모에서 인스턴스를 평가합니다. TTT-Linear와 TTT-MLP는 모두 기준선과 일치하거나 초과합니다. Transformer와 유사하게 더 많은 토큰을 조건으로 하여 복잡도를 계속 줄일 수 있는 반면 Mamba는 16k 맥락 이후에는 그럴 수 없습니다. 예비 시스템 최적화를 통해 TTT-Linear는 이미 8k 컨텍스트에서 Transformer보다 빠르고 벽시계 시간에서 Mamba와 일치합니다. TTT-MLP는 여전히 메모리 I/O에서 문제에 직면해 있지만 긴 컨텍스트에서 더 큰 잠재력을 보여주며 미래 연구의 유망한 방향을 가리킵니다.

 

배경 및 문제점
기존의 RNN과 LSTM(Long Short-Term Memory)은 긴 문맥(context)을 처리하는 데 한계가 있으며, 이는 특히 2020년에 OpenAI에서 발표한 연구에서 강조되었습니다. RNN은 입력 시퀀스를 고정 크기의 숨겨진 상태로 압축해야 하기 때문에 긴 문맥을 효과적으로 활용하지 못합니다.

TTT 레이어 소개
TTT 레이어는 테스트 시퀀스에서 모델을 업데이트하는 새로운 유형의 시퀀스 모델링 레이어입니다. 이 레이어의 숨겨진 상태는 모델 자체이며, 업데이트 규칙은 자가 지도 학습(self-supervised learning)의 한 단계입니다. TTT 레이어는 RNN 레이어나 자기 주의(self-attention)와 같은 다른 네트워크 아키텍처에 통합될 수 있습니다.

TTT-Linear 및 TTT-MLP
TTT 레이어의 두 가지 간단한 구현 방식인 TTT-Linear와 TTT-MLP가 제안되었습니다. 이들 방식은 각각 선형 모델과 2층 MLP(Multi-Layer Perceptron)를 숨겨진 상태로 사용합니다. 실험 결과, TTT-Linear는 Transformer와 Mamba보다 더 나은 성능을 보여주었습니다.

성능 향상
TTT 레이어는 플롭(FLOPs) 측면에서 효율적이며, 두 가지 혁신적인 방법을 통해 실제 시간에서도 효율적입니다. 첫 번째는 미니배치(mini-batch)를 사용한 병렬화이고, 두 번째는 듀얼 폼(dual form)을 사용하여 현대 GPU와 TPU의 성능을 극대화하는 것입니다.

실험 결과
TTT-Linear는 Transformer 및 Mamba와 비교하여 긴 문맥에서도 더 낮은 퍼플렉시티(perplexity)를 유지할 수 있습니다. 또한, TTT 레이어는 모델의 크기가 커질수록 더 큰 장점을 나타냅니다.

향후 연구 방향
문서에서는 TTT 레이어의 잠재력을 극대화하기 위한 여러 가지 향후 연구 방향을 제안하고 있습니다. 여기에는 시스템 최적화, 더 긴 문맥 및 더 큰 모델, 더 야심 찬 f의 구현 등이 포함됩니다.

이 문서의 연구는 인간 지능의 본질을 탐구하는 AI 연구의 일환으로, 인간 학습과 TTT의 연결 가능성을 모색합니다.

TTT는 모델이 테스트 데이터를 처리하면서도 스스로 학습을 계속하도록 설계된 방법입니다. 이를 통해 모델은 새로운 데이터에 대해 더 적응력 있게 되며, 기존의 학습 방식보다 더 나은 성능을 발휘할 수 있습니다.

TTT의 주요 아이디어는 모델의 숨겨진 상태를 단순한 값이 아닌 학습 가능한 모델로 만드는 것입니다. 테스트 데이터가 들어올 때마다 이 숨겨진 모델은 자가 지도 학습(self-supervised learning)을 통해 업데이트됩니다.

주요 특징
자가 지도 학습(Self-Supervised Learning): TTT에서는 숨겨진 상태를 업데이트하는 규칙이 자가 지도 학습의 단계로 이루어져 있습니다.

모델의 두 가지 구현 방식:
TTT-Linear: 숨겨진 상태가 선형 모델인 경우.
TTT-MLP: 숨겨진 상태가 두 층의 MLP인 경우.
효율성: TTT-Linear는 FLOP(부동 소수점 연산) 측면에서 Transformer보다 효율적이며, 긴 문맥에서도우수한 성능을 보입니다.

미니배치 및 듀얼 폼: TTT는 미니배치와 듀얼 폼을 사용하여 병렬화와 현대 GPU/TPU의 성능을 극대화합니다.

장점
적응력 향상: 테스트 데이터를 처리하면서 모델이 실시간으로 학습하여 더 나은 예측 성능을 발휘할 수 있습니다.
긴 문맥 처리: TTT는 긴 문맥에서도 효율적으로 작동하며, 더 많은 정보를 활용하여 성능을 향상시킬 수 있습니다.

실시간 학습 예시
예를 들어, TTT-Linear 모델은 테스트 시퀀스를 처리할 때 각 입력 토큰에 대해 숨겨진 상태를 업데이트하고, 이를 통해 점진적으로 더 나은 예측을 수행할 수 있습니다. 이 과정은 일반적인 모델 훈련 과정과 유사하지만, 테스트 중에도 학습이 이루어진다는 점에서 차이가 있습니다.

따라서 TTT는 실시간 학습이 가능하며, 특히 새로운 데이터에 빠르게 적응해야 하는 상황에서 유용합니다. 자세한 기술적 구현과 성능 비교는 원문을 참조하면 도움이 될 것입니다.