언어 모델을 완전히 변경합니다. 새로운 아키텍처 TTT가 Transformer를 능가하고 ML 모델이 RNN 숨겨진 상태를 대체합니다.

대형 모델의 성능이 125M에서 1.3B로 향상되었습니다.

믿을 수 없는 일이 드디어 일어났습니다.

지금까지 AI 분야에서 인기를 끌었던 Transformer를 대체할 새로운 LLM(대형 언어 모델) 아키텍처가 예상되며 성능은 Mamba보다 뛰어납니다. 월요일 인공지능 커뮤니티에서는 TTT(Test-Time Training)에 관한 논문이 화제가 되었습니다.

논문 링크: https://arxiv.org/abs/2407.04620

이 연구의 저자는 스탠포드 대학, 캘리포니아 대학, 버클리, 캘리포니아 대학, 샌디에고 및 Meta 출신입니다. 그들은 RNN의 숨겨진 상태를 기계 학습 모델로 대체하는 새로운 아키텍처 TTT를 설계했습니다. 모델은 입력 토큰의 실제 경사하강법을 통해 컨텍스트를 압축합니다.

이 연구의 저자 중 한 명인 Karan Dalal은 이것이 언어 모델 접근 방식을 근본적으로 바꿀 것이라고 믿습니다.

기계 학습 모델에서 TTT 계층은 Attention을 직접 대체하고 표현 메모리를 통해 선형 복잡성 아키텍처를 잠금 해제하여 상황에 맞는 수백만(때로는 수십억) 개의 토큰으로 LLM을 교육할 수 있습니다.

저자는 매개변수 크기가 125M에서 1.3B까지인 대형 모델에 대해 일련의 비교를 수행한 결과 TTT-Linear와 TTT-MLP가 모두 가장 강력한 Transformer 및 Mamba 아키텍처 방법과 일치하거나 패배할 수 있음을 발견했습니다.

새로운 정보 압축 및 모델 메모리 메커니즘인 TTT 레이어는 Transformer의 self-attention 레이어를 간단하고 직접적으로 대체할 수 있습니다.

Mamba와 비교하여 TTT-Linear는 복잡성이 낮고 FLOP가 적으며(왼쪽) 긴 컨텍스트를 더 잘 활용합니다(오른쪽).

이것은 이론 복잡성이 선형일 뿐만 아니라 실제 실행에서도 시간도 빨라지고.

• 논문이 온라인에 게재된 후 저자는 사람들이 훈련하고 테스트할 수 있도록 코드와 jax를 공개했습니다: https://github.com/test-time-training/ttt-lm-jax
• 또한 PyTorch 추론 코드: https://github.com/test-time-training/ttt-lm-pytorch

방법 소개

긴 컨텍스트에 대한 과제는 RNN 레이어의 특성: self-attention 메커니즘과 달리 RNN 레이어는 컨텍스트를 고정된 크기의 숨겨진 상태로 압축해야 하며 업데이트 규칙은 수천 또는 심지어 수백만 개의 토큰 간의 기본 구조와 관계를 발견해야 합니다.

연구팀은 먼저 자기 지도 학습이 대규모 훈련 세트를 LLM과 같은 모델의 가중치로 압축할 수 있다는 점을 관찰했으며, LLM 모델은 종종 훈련 데이터 간의 의미론적 연결에 대한 깊은 이해를 보여줍니다.

이 관찰에서 영감을 받아 연구팀은 숨겨진 상태가 모델이고 업데이트 규칙이 자기 지도 학습 단계인 새로운 클래스의 시퀀스 모델링 레이어를 설계했습니다. 테스트 시퀀스의 숨겨진 상태를 업데이트하는 프로세스는 테스트 시 모델을 교육하는 것과 동일하므로 연구팀은 이 새로운 레이어를 TTT(Test-Time Training) 레이어라고 부릅니다.

연구팀은 TTT-Linear와 TTT-MLP라는 두 가지 간단한 예를 소개합니다. 여기서 숨겨진 상태는 각각 선형 모델과 2계층 MLP입니다. TTT 레이어는 RNN 레이어 및 self-attention과 유사하게 모든 네트워크 아키텍처에 통합되고 엔드투엔드를 최적화할 수 있습니다.

TTT 레이어를 더 효율적으로 만들기 위해 연구에서는 TTT 레이어를 개선하기 위한 몇 가지 트릭을 채택했습니다.

첫 번째, 더 나은 병렬성을 얻기 위해 정규 학습 중에 미니 배치 시퀀스에 대한 그라데이션 단계를 취하는 것과 유사하게, 연구에서는 TTT 동안 소량의 토큰을 사용하세요.

두 번째로, 이 연구는 최신 GPU 및 TPU를 더 잘 활용하기 위해 각 TTT 미니 배치 내에서 작업을 위한 이중 형식을 개발합니다. 이중 형식의 출력은 단순 구현과 동일하지만 훈련 속도가 5배 이상 빠릅니다. 그림 3에서 볼 수 있듯이 TTT-Linear는 Transformer보다 빠르며 8k 환경에서 Mamba와 비슷합니다.

연구팀은 모든 시퀀스 모델링 레이어가 그림 4와 같이 역사적 맥락을 숨겨진 상태로 저장하는 것으로 볼 수 있다고 믿습니다.

예를 들어 LSTM, RWKV 및 Mamba 레이어와 같은 RNN 레이어는 시간이 지남에 따라 컨텍스트를 고정 크기 상태로 압축합니다. 이 압축에는 두 가지 결과가 있습니다. 한편으로는 각 토큰에 대한 업데이트 규칙과 출력 규칙에 일정한 시간이 필요하기 때문에 입력 토큰 x_t를 출력 토큰 z_t에 매핑하는 것이 효율적입니다. 반면, 긴 컨텍스트에서 RNN 레이어의 성능은 숨겨진 상태 s_t의 표현력에 의해 제한됩니다.

Self-attention은 숨겨진 상태(키-값 캐시라고도 함)가 t에 따라 선형적으로 증가하는 목록이라는 점을 제외하면 위의 관점에서도 볼 수 있습니다. 업데이트 규칙은 단순히 현재 KV 튜플을 이 목록에 추가하는 반면, 출력 규칙은 t 이전의 모든 튜플을 검색하여 어텐션 매트릭스를 형성합니다. 숨겨진 상태는 압축 없이 모든 기록 컨텍스트를 명시적으로 저장하므로 긴 컨텍스트에 대한 RNN 레이어보다 self-attention이 더 표현력이 좋습니다. 그러나 선형적으로 증가하는 숨겨진 상태를 스캔하는 데 필요한 시간도 선형적으로 증가합니다. 긴 컨텍스트를 효율적이고 표현력 있게 유지하려면 연구자에게는 더 나은 압축 휴리스틱이 필요합니다. 특히, 수천 또는 수백만 개의 토큰을 기본 구조와 관계를 효과적으로 포착하는 숨겨진 상태로 압축해야 합니다. 어렵게 들릴 수도 있지만 실제로 많은 사람들이 이 경험적 방법에 매우 익숙합니다.

백본 아키텍처. RNN 레이어를 더 큰 아키텍처에 통합하는 가장 깔끔한 방법은 여기서 백본이라고 불리는 Transformer의 셀프 어텐션을 직접 교체하는 것입니다. 그러나 기존 RNN(예: Mamba 및 Griffin)은 Transformer와 다른 백본 레이어를 사용합니다. 특히 백본 레이어에는 RNN 레이어 앞에 시간적 컨볼루션이 포함되어 있어 시간에 따른 로컬 정보를 수집하는 데 도움이 될 수 있습니다. 연구진은 Mamba 백본을 실험한 결과 TTT 레이어의 Perplexity도 개선할 수 있다는 사실을 발견하여 그림 16과 같이 제안된 방법에 포함시켰습니다.

실험 결과

실험에서 연구자들은 TTT-Linear와 TTT-MLP를 Transformer와 Mamba, 두 가지 기준선과 비교했습니다.

짧은 텍스트

그림 11에서 다음과 같은 결론을 내릴 수 있습니다.

• 2k 컨텍스트에서 TTT-Linear(M), Mamba 및 Transformer의 성능은 비슷합니다. 선의 대부분이 겹칩니다. TTT-MLP(M)는 FLOP 예산이 클수록 성능이 약간 떨어집니다. TTT-MLP는 다양한 모델 크기에서 TTT-Linear보다 복잡도가 더 높지만 FLOP의 추가 비용이 이러한 이점을 상쇄합니다.
• 8k 컨텍스트의 경우 TTT-Linear(M)와 TTT-MLP(M) 모두 Mamba보다 성능이 훨씬 뛰어나며 이는 2k 컨텍스트의 관찰과는 상당히 다릅니다. Transformer 백본 네트워크를 사용하는 TTT-MLP(T)도 약 1.3B로 Mamba보다 약간 더 좋습니다. 중요한 현상은 컨텍스트 길이가 증가함에 따라 Mamba 계층에 비해 TTT 계층의 장점도 확장된다는 것입니다.
• 컨텍스트 길이가 8k에 도달하면 Transformer는 각 모델 크기에서 여전히 우수한 성능을 발휘하지만 FLOP 비용으로 인해 더 이상 경쟁력이 없습니다.

위 결과는 TTT 레이어를 Mamba 백본 네트워크에서 Transformer 백본 네트워크로 전환할 때의 영향을 보여줍니다. 연구원들은 시퀀스 모델링 계층의 숨겨진 상태가 덜 표현될 때 Mamba 백본 네트워크의 시간적 컨볼루션이 더 유용하다는 가설을 세웠습니다. 선형 모델은 MLP보다 표현력이 떨어지므로 컨볼루션의 이점을 더 많이 얻습니다.

긴 텍스트: 책

긴 컨텍스트의 기능을 평가하기 위해 Pile의 인기 있는 하위 집합인 Books3을 사용하여 1k에서 32k까지 2x 증분으로 컨텍스트 길이를 실험했습니다. 여기서의 훈련 방법은 Pile과 동일하며 TTT 레이어에 대한 모든 실험은 한 번의 훈련 실행으로 수행됩니다. 그림 12의 결과 하위 집합에서 그들은 다음과 같은 관찰을 했습니다.

Books 2k의 맥락에서 Pile 2k에 대한 모든 관찰은 여전히 ​​유지됩니다. 단, Mamba는 이제 TTT-Linear보다 약간 더 나은 성능을 발휘합니다. Pile 2k에서는 선이 대략 겹칩니다.

32k 맥락에서 TTT-Linear(M)와 TTT-MLP(M)는 모두 Pile 8k에 대한 관찰과 유사하게 Mamba보다 더 나은 성능을 발휘합니다. Transformer 백본을 갖춘 TTT-MLP(T)도 32k 컨텍스트에서는 Mamba보다 성능이 약간 더 좋습니다.

TTT-MLP(T)는 1.3B 규모에서 TTT-MLP(M)보다 약간 더 나쁠 뿐입니다. 위에서 언급한 바와 같이, 명확한 선형 적합도가 부족하여 경험적 스케일링 법칙을 도출하기가 어렵습니다. 그러나 TTT-MLP(T)의 강력한 추세는 Transformer 백본이 우리 평가 범위를 넘어서 더 큰 모델과 더 긴 컨텍스트에 더 적합할 수 있음을 시사합니다.

시계 시간

LLM의 훈련과 추론은 순방향, 역방향, 세대로 분해될 수 있습니다. 추론 중 큐 워드 처리(사전 채우기라고도 함)는 역방향 작업에 중간 활성화 값을 저장할 필요가 없다는 점을 제외하면 훈련 중 순방향 작업과 동일합니다.

순방향(훈련 및 추론 중)과 역방향이 모두 병렬로 처리될 수 있으므로 여기서는 이중 형식이 사용됩니다. 새로운 토큰 생성(디코딩이라고도 함)은 본질적으로 순차적이므로 여기서는 원시 형식이 사용됩니다.

연구원은 리소스 제한으로 인해 이 기사의 실험은 JAX로 작성되었으며 TPU에서 실행되었다고 언급했습니다. v5e-256 TPU Pod에서 Transformer 기준은 2k 컨텍스트로 학습하는 데 반복당 0.30초가 걸리는 반면, TTT-Linear는 반복당 0.27초가 소요됩니다. 이는 시스템 최적화 없이 10% 더 빠릅니다. Mamba(PyTorch, Triton 및 CUDA로 구현됨)는 GPU에서만 실행될 수 있으므로 공정한 비교를 위해 연구원들은 이 방법에 대해 예비 시스템 최적화를 수행하여 GPU에서 실행될 수 있도록 했습니다.

그림 15의 왼쪽은 배치 크기 16에서 각 모델에 대한 순방향 커널의 대기 시간을 보여줍니다. 모든 모델은 1.3B입니다(Mamba는 1.4B). 여기서는 HuggingFace Transformer 대신 vLLM이 사용되었기 때문에 여기의 Transformer 기준선이 Mamba 논문의 기준선보다 훨씬 빠르다는 점은 주목할 가치가 있습니다.

또한 연구원들은 세대를 위한 또 다른 GPU 커널을 작성하고 그림 15의 오른쪽에 배치 크기 512로 속도를 벤치마킹했습니다. 일반적으로 사용되는 또 다른 벽시계 시간 측정 기준은 처리량입니다. 이는 더 큰 배치 크기를 사용할 때의 잠재적 이점을 고려합니다. 처리량의 경우 위의 모든 관찰과 방법 간 순서가 여전히 유지됩니다.

주저자

TTT 연구가 제출된 후 논문의 저자 중 한 명인 UCSD 조교수 Xiaolong Wang이 축하 메시지를 트윗했습니다. 그는 TTT에 대한 연구는 1년 반 동안 진행됐지만 실제로는 TTT(Test Time Training)라는 아이디어가 탄생한 지 5년이 됐다고 말했다. 원래 아이디어와 현재 결과는 완전히 다르지만.

TTT 논문의 세 명의 주요 저자는 각각 Stanford, UC Berkeley 및 UCSD 출신입니다.

그 중 유선님은 스탠포드 대학교에서 박사후 연구원으로 UC Berkeley EECS에서 박사 학위를 취득했으며, 장기 연구 방향은 TTT입니다.

Xinhao Li는 UCSD에서 박사 과정을 밟고 있으며 중국 전자 과학 기술 대학교를 졸업했습니다.

Karan Dalal은 UC Berkeley의 박사 과정 지원자이며 고등학교 시절 Otto라는 수의학 원격 의료 스타트업을 공동 창립했습니다.

위 세 사람은 모두 개인 홈페이지 첫 줄에 연구 방향을 소개하는 시험 훈련 글을 썼습니다.

자세한 연구 내용은 원문을 참고해주세요.

위 내용은 언어 모델을 완전히 변경합니다. 새로운 아키텍처 TTT가 Transformer를 능가하고 ML 모델이 RNN 숨겨진 상태를 대체합니다.의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!