Learning to (Learn at Test Time): RNNswith Expressive Hidden States

https://arxiv.org/abs/2407.04620v1

Learning to (Learn at Test Time): RNNs with Expressive Hidden States

 

요약

자기-주의(Self-attention)는 긴 문맥에서 뛰어난 성능을 보이지만, 이차 복잡성을 가지고 있습니다. 기존의 RNN 레이어는 선형 복잡성을 가지지만, 그들의 숨겨진 상태의 표현력 한계로 긴 문맥에서의 성능이 제한됩니다. 우리는 선형 복잡성과 표현력이 높은 숨겨진 상태를 가진 새로운 유형의 시퀀스 모델링 레이어를 제안합니다. 핵심 아이디어는 숨겨진 상태 자체를 기계 학습 모델로 만들고, 업데이트 규칙을 자기 지도 학습의 한 단계로 만드는 것입니다. 숨겨진 상태가 테스트 시퀀스에서도 훈련을 통해 업데이트되므로, 우리의 레이어는 테스트-시간 훈련(TTT) 레이어라 불립니다. 우리는 TTT-Linear와 TTT-MLP라는 두 가지 구현을 고려하며, 이들의 숨겨진 상태는 각각 선형 모델과 2층 MLP입니다. 우리는 125M에서 1.3B 파라미터 규모에서 강력한 Transformer 및 최신 RNN인 Mamba와 비교하여 우리의 구현을 평가합니다. TTT-Linear와 TTT-MLP는 기준선과 일치하거나 초과합니다. Transformer와 유사하게, 더 많은 토큰에 조건부로 할당함으로써 당혹감을 계속 줄일 수 있지만, Mamba는 16k 문맥 이후로는 그렇지 못합니다. 초기 시스템 최적화로, TTT-Linear는 이미 8k 문맥에서 Transformer보다 빠르며, Mamba와 벽-시계 시간에서 일치합니다. TTT-MLP는 여전히 메모리 I/O에서 도전에 직면하고 있지만, 긴 문맥에서 더 큰 잠재력을 보여 미래 연구의 유망한 방향을 제시합니다.

 

-----------------------------------------------------------------------------------------------------------------------------------------------------------------

당혹감 = perplexity

Perplexity는 주로 자연어 처리(NLP)에서 사용되는 평가 척도 중 하나입니다. 언어 모델의 성능을 측정하는 데 사용되며, 모델이 주어진 텍스트 데이터를 얼마나 잘 예측하는지를 나타냅니다. Perplexity 값이 낮을수록 모델이 더 잘 예측하는 것을 의미합니다.

-----------------------------------------------------------------------------------------------------------------------------------------------------------------

 

그림 1

모든 시퀀스 모델링 레이어는 업데이트 규칙에 따라 전환되는 숨겨진 상태로 표현될 수 있습니다. 우리의 핵심 아이디어는 숨겨진 상태 자체를 가중치 𝑊 를 가진 모델 𝑓 로 만들고, 업데이트 규칙을 자기 지도 손실 ℓ 에 대한 기울기 단계로 만드는 것입니다. 따라서 테스트 시퀀스에서 숨겨진 상태를 업데이트하는 것은 테스트 시간에 모델 𝑓 를 훈련하는 것과 동일합니다. 이 과정은 테스트-시간 훈련(TTT)으로 알려져 있으며, 우리의 TTT 레이어에 프로그램되어 있습니다.

 

그림 2

Mamba와 비교할 때, TTT-Linear는 당혹감과 FLOP가 더 적으며(왼쪽), 긴 문맥을 더 잘 활용합니다(오른쪽). 평가들은 Kaplan et al. [40]을 따릅니다. 왼쪽: 350M에서 1.3B 파라미터 사이에서 책 스케일링 추세를 확대했습니다. 760M 및 1.3B에서 TTT-Linear는 Mamba보다 적은 FLOP로 당혹감을 초과하며, 선형 보간 하에서 Transformer를 초과합니다. 오른쪽: Transformer와 TTT-Linear는 더 많은 토큰에 조건부로 할당함에 따라 당혹감을 계속 줄일 수 있지만, Mamba는 16k 문맥 이후로는 그렇지 못합니다. 모든 방법은 Mamba 1.4B와 동일한 훈련 FLOP를 가지고 있습니다. 세부 사항은 3.2절에 있습니다.

 

그림 3

배치 크기 16에서 문맥 길이가 변함에 따라 토큰당 전달 시간(지연 시간). 모든 모델은 1.3B(1.4B는 Mamba). 문맥 길이가 증가함에 따라 Transformer의 토큰당 전달 시간은 선형으로 증가하지만, 다른 두 방법은 거의 일정하게 유지됩니다. TTT-Linear는 8k 문맥에서 Transformer보다 빠르며 Mamba와 일치합니다. 세부 사항은 3.3절에 있습니다.

 

1. 서론

2020년, OpenAI의 확장 법칙 논문(Kaplan et al. [40])은 LSTM(일종의 RNN)이 Transformer와 같은 방식으로 확장되지 않으며 긴 문맥을 효과적으로 사용하지 못한다고 밝혔습니다. 이제 최신 RNN과 최선의 실천 방안을 통해, 우리는 그림 2에서 이러한 결과를 재평가합니다.

 

왼쪽 그림에서는 오늘날 가장 인기 있는 RNN 중 하나인 Mamba [26]가 강력한 Transformer와 유사하게 확장된다는 것을 알 수 있습니다. 이는 2020년의 LSTM 이후 큰 진전을 보인 것입니다. 그러나 오른쪽 그림에서는 Mamba가 LSTM과 같은 문제를 겪고 있음을 알 수 있습니다. 시퀀스의 후반부 토큰은 더 많은 정보를 조건으로 하기 때문에 평균적으로 예측하기가 더 쉬워야 합니다. 이는 32k 문맥에서 각 토큰 인덱스에서 평균 당혹감이 감소하는 Transformer의 경우에 해당합니다. 반면 Mamba는 16k 문맥 이후에 같은 지표가 정체됩니다.

 

이 결과는 기존 RNN의 불편한 현실을 나타냅니다. 한편으로, RNN의 주요 장점(Transformer와 비교하여)은 선형(대신 이차) 복잡성입니다. 이 점근적 장점은 실질적으로는 긴 문맥에서만 실현되며, 그림 3에 따르면 이는 8k 이후입니다. 그러나 문맥이 충분히 길어지면, Mamba와 같은 기존 RNN은 조건부로 할당되는 추가 정보를 실제로 활용하는 데 어려움을 겪습니다.

 

긴 문맥과의 어려움은 RNN 레이어의 본질적인 특성에서 기인합니다. 자기-주의와 달리, RNN 레이어는 문맥을 고정된 크기의 숨겨진 상태로 압축해야 합니다. 압축 휴리스틱으로서, 업데이트 규칙은 수천 또는 수백만 개의 토큰 간의 기본 구조와 관계를 발견해야 합니다. 이 논문에서는 자기 지도 학습이 대규모 훈련 세트를 LLM과 같은 모델의 가중치로 압축할 수 있으며, 이는 훈련 데이터 간의 의미적 연결에 대한 깊은 이해를 자주 보여주며, 우리가 압축 휴리스틱에서 필요로 하는 것과 정확히 일치하는 것을 관찰합니다.

TTT 레이어

이 관찰에 영감을 받아, 숨겨진 상태가 모델이고 업데이트 규칙이 자기 지도 학습의 단계인 새로운 유형의 시퀀스 모델링 레이어를 설계합니다. 테스트 시퀀스에서 숨겨진 상태를 업데이트하는 과정이 테스트 시간에 모델을 훈련하는 것과 동일하기 때문에, 이 새로운 유형의 레이어를 테스트-시간 훈련(TTT) 레이어라고 합니다. 우리는 이 클래스 내에서 두 가지 간단한 구현을 소개합니다: 숨겨진 상태가 선형 모델인 TTT-Linear와 두 층 MLP인 TTT-MLP. TTT 레이어는 RNN 레이어와 자기-주의와 유사하게 네트워크 아키텍처에 통합되고 엔드 투 엔드로 최적화될 수 있습니다.

실제 시간 효율성

TTT 레이어는 이미 FLOP에서 효율적이지만, 실제 시간에서도 효율적이도록 두 가지 실용적인 혁신을 제안합니다. 첫째, 더 나은 병렬 처리를 위해 정규 훈련 동안 시퀀스의 미니 배치에 대한 기울기 단계를 취하는 표준 관행과 유사하게, 우리는 TTT 동안 토큰의 미니 배치를 사용합니다. 둘째, 각 TTT 미니 배치 내의 작업을 더 잘 활용하기 위해 현대 GPU와 TPU의 이점을 더 잘 활용하기 위해 이중 형식을 개발합니다. 이중 형식은 출력에서 순진한 구현과 동일하지만 훈련 속도는 5배 이상 빠릅니다. 그림 3에서 볼 수 있듯이, TTT-Linear는 8k 문맥에서 Transformer보다 빠르고 Mamba와 일치합니다.

평가 및 해결 과제

논문 초반에 TTT-Linear의 일부 결과를 강조했지만, 3장에서 TTT-Linear와 TTT-MLP에 대한 더 포괄적인 평가와 우리의 평가로 드러난 해결 과제를 제시합니다. 예를 들어, Chinchilla 레시피 [34]를 따르는 우리의 평가에서는 Transformer 기준선조차 선형 확장 추세에 깔끔하게 맞지 않습니다. 우리의 학술 자원에 의해 제한되어, 우리는 커뮤니티가 이 문제에 대한 해결책을 탐구하는 데 함께하길 권장합니다.

기여 요약

1. 우리는 숨겨진 상태가 모델이고 업데이트 규칙이 자기 지도 학습인 새로운 유형의 시퀀스 모델링 레이어인 TTT 레이어를 제안합니다. 레이어의 순방향 통과가 자체적으로 훈련 루프를 포함한다는 우리의 관점은 미래 연구의 새로운 방향을 열어줍니다.
2. TTT 레이어의 간단한 구현인 TTT-Linear는 125M에서 1.3B 파라미터 범위의 평가에서 Transformer와 Mamba를 능가합니다.
3. 미니 배치 TTT와 이중 형식을 통해 TTT 레이어의 하드웨어 효율성을 향상시켜, TTT-Linear를 LLM의 실용적인 구성 요소로 만듭니다.

2. 방법

모든 시퀀스 모델링 레이어는 과거 문맥을 숨겨진 상태에 저장하는 관점에서 볼 수 있습니다. 그림 4에서 보듯이, RNN 레이어(LSTM [33], RWKV [56], Mamba [26] 등)는 시간 경과에 따라 문맥을 고정 크기의 상태로 압축합니다. 이 압축에는 두 가지 결과가 있습니다. 첫째, 입력 토큰 𝑥 𝑡 을 출력 토큰 𝑧 𝑡 으로 매핑하는 것이 효율적입니다. 이는 업데이트 규칙과 출력 규칙이 토큰당 일정한 시간이 걸리기 때문입니다. 둘째, 긴 문맥에서 RNN 레이어의 성능은 숨겨진 상태 𝑠 𝑡 의 표현력에 의해 제한됩니다.

 

자기-주의(Self-attention)도 위의 관점에서 볼 수 있지만, 자기-주의의 숨겨진 상태는 키-값(KV) 캐시로 알려진 리스트로, 𝑡 에 따라 선형적으로 증가합니다. 업데이트 규칙은 현재 KV 튜플을 이 리스트에 추가하고, 출력 규칙은 𝑡 까지의 모든 튜플을 스캔하여 어텐션 매트릭스를 형성합니다. 숨겨진 상태는 모든 과거 문맥을 압축 없이 명시적으로 저장하므로, 자기-주의는 긴 문맥에서 RNN 레이어보다 표현력이 더 높습니다. 그러나 이 선형적으로 증가하는 숨겨진 상태를 스캔하는 데에도 토큰당 선형적으로 시간이 증가합니다.

 

긴 문맥에서 효율적이면서도 표현력이 높으려면 더 나은 압축 휴리스틱이 필요합니다. 구체적으로, 수천 또는 수백만 개의 토큰을 숨겨진 상태로 압축하여 이들의 기본 구조와 관계를 효과적으로 캡처할 수 있어야 합니다. 이는 어려운 작업처럼 들리지만, 사실 우리 모두는 이미 이런 휴리스틱에 익숙합니다.

 

그림 4.1 설명

그림 4은 시퀀스 모델링 레이어가 숨겨진 상태를 통해 과거 문맥을 저장하는 방식을 보여줍니다. RNN 레이어는 시간 경과에 따라 문맥을 고정 크기의 상태로 압축하고, 자기-주의는 키-값 캐시를 통해 모든 과거 문맥을 압축 없이 저장합니다. 이 두 접근법은 각각의 장단점이 있으며, 긴 문맥에서 효율적이면서도 표현력이 높은 모델을 만들기 위해서는 새로운 압축 휴리스틱이 필요합니다.

 

 

 

 

class TTT_Layer(nn.Module):
    def __init__(self):
        self.task = Task()

    def forward(self, in_seq):
        state = Learner(self.task)
        out_seq = []
        for tok in in_seq:
            state.train(tok)
            out_seq.append(state.predict(tok))
        return out_seq

class Task(nn.Module):
    def __init__(self):
        self.theta_K = nn.Parameter((d1, d2))
        self.theta_V = nn.Parameter((d1, d2))
        self.theta_Q = nn.Parameter((d1, d2))

    def loss(self, f, x):
        train_view = self.theta_K @ x
        label_view = self.theta_V @ x
        return MSE(f(train_view), label_view)

class Learner():
    def __init__(self, task):
        self.task = task
        # 선형 모델이지만, 다른 모델도 사용할 수 있음
        self.model = Linear()
        # 단순화를 위해 온라인 기울기 하강 사용
        self.optim = OGD()

    def train(self, x):
        # 손실의 첫 번째 인수에 대한 기울기 함수
        grad_fn = grad(self.task.loss)
        # 내부 루프 기울기 계산
        grad_in = grad_fn(self.model, x)
        # 현재 파라미터에서 시작하여 grad_in 방향으로 한 단계 이동
        self.optim.step(self.model, grad_in)

    def predict(self, x):
        test_view = self.task.theta_Q @ x
        return self.model(test_view)

그림 6

PyTorch 스타일로 선형 모델과 온라인 기울기 하강을 사용하는 TTT 레이어의 순진한 구현. TTT_Layer는 다른 시퀀스 모델링 레이어와 마찬가지로 더 큰 네트워크에 쉽게 삽입될 수 있습니다. 네트워크를 훈련시키면 TTT_Layer의 Task 파라미터가 최적화됩니다. Learner는 nn.Module의 하위 클래스가 아니므로, 각 state.train 호출 시 내부 루프에서 state.model이 수동으로 업데이트됩니다. 간단하게 하기 위해, 가끔 model을 model.parameters로 과부하할 수 있습니다.

 

-----------------------------------------------------------------------------------------------------------------------------------------------------------------

TTT_Layer는 시퀀스의 각 토큰을 처리하며, Task 클래스는 파라미터를 초기화하고 손실 함수를 정의합니다. Learner 클래스는 Task 객체를 기반으로 학습과 예측을 수행합니다.

-----------------------------------------------------------------------------------------------------------------------------------------------------------------

 

 

그림 7: 첫 번째 TTT 미니배치의 고수준 계산 그래프

노드는 변수이고, 엣지는 계산입니다. 파란색 노드는 입력 변수이고, 노란색은 출력 변수입니다.

 

 

2.6 이론적 등가성

 

그림 9

매개변수 학습자는 두 가지 속성을 정의해야 합니다: 모델과 옵티마이저(왼쪽). 각 학습자는 고유하게 TTT 레이어를 유도하며(오른쪽), 본 논문에서는 TTT-Linear와 TTT-MLP라는 두 가지 유도된 TTT 레이어를 제안합니다. 선형 모델과 배치 GD를 사용하는 TTT 레이어는 선형 어텐션 [41]과 동등합니다.

 

실험적 검증

표 1에서 우리는 선형 어텐션의 개선된 구현을 통해 위의 동등성을 실험적으로 검증합니다. 선형 어텐션의 원래 공식은 정규화와 x_t​에 대한 특징 확장을 포함하며, 이는 동등한 TTT 레이어에 여전히 포함될 수 있습니다. 그러나 이전 연구에 따르면, 이러한 두 가지 추가는 성능에 악영향을 미칠 수 있으며, 이는 우리 실험에서도 확인되었습니다(표 1의 첫 번째와 두 번째 행 비교). 따라서 우리는 두 가지 추가 사항 없이 선형 어텐션의 가장 간단한 공식에 해당하는 TTT 레이어를 구성합니다. 그런 다음, 각 구성 요소의 기여도를 설명하기 위해 TTT 레이어에 하나씩 추가하여 최종적으로 TTT-Linear라는 제안된 인스턴스를 얻습니다. 배치 GD에서 미니배치 GD로의 변경이 가장 큰 개선을 가져왔습니다.

 

그림 10: RNN 레이어와 TTT 레이어

그림 10에서는 RNN 레이어와 TTT 레이어가 모두 시퀀스 모델링 레이어의 하위 집합임을 보여줍니다. RNN 레이어는 시간이 지나도 크기가 고정된 숨겨진 상태를 가집니다. 매개변수 학습자를 가진 TTT 레이어도 RNN 레이어로 간주될 수 있습니다. 이는 해당 레이어의 숨겨진 상태도 고정 크기이기 때문입니다. 비매개변수 학습자를 가진 TTT 레이어는 자기-어텐션을 나타낼 수 있으며, 이는 2.6절에서 논의되었습니다.

 

 

 

 

3 실험

우리는 TTT-Linear와 TTT-MLP를 Transformer와 현대적인 RNN인 Mamba와 비교하여 평가합니다. 우리의 주요 코드베이스는 JAX에서 LLM을 훈련하고 제공하기 위한 오픈소스 프로젝트인 EasyLM [25]를 기반으로 합니다. 모든 실험은 첫 페이지 하단에 제공된 공개 코드와 데이터셋을 사용하여 재현할 수 있습니다.

데이터셋

Mamba 논문 [26]을 따라, 우리는 Pile [24]에서 2k 및 8k 문맥 길이를 사용하여 표준 실험을 수행합니다. Pile은 오픈소스 LLM을 훈련하기 위한 문서 데이터셋으로 인기가 높습니다 [9]. 그러나 Pile에는 8k보다 긴 시퀀스가 거의 없습니다 [19]. 긴 문맥에서의 성능을 평가하기 위해, 우리는 Pile의 하위 집합인 Books3를 사용하여 1k에서 32k까지의 문맥 길이에서 실험을 수행합니다. 이 데이터셋은 긴 문맥에서 LLM을 훈련하는 데 널리 사용되었습니다 [49, 3].

백본 아키텍처

2.7절에서 논의한 바와 같이, Transformer와 Mamba는 서로 다른 백본을 사용하며, 특별한 언급이 없는 한 TTT-Linear와 TTT-MLP는 항상 Mamba 백본을 사용합니다. 소거 실험으로, 그림 11과 그림 12는 Transformer 백본 내에서 TTT 레이어를 포함합니다. 한 그림에 Transformer 백본과 Mamba 백본이 모두 포함된 경우, 각각 (T)와 (M)으로 표시합니다.

프로토콜

기준선에 대한 공정성을 보장하기 위해, 우리는 가능한 한 Mamba 논문의 평가 프로토콜을 엄격히 따릅니다.

• 각 평가 설정(예: 데이터셋, 문맥 길이 및 방법)마다 네 가지 모델 크기(125M, 350M, 760M, 1.3B 파라미터)로 실험합니다. Mamba의 경우, 해당 크기는 130M, 370M, 790M, 1.4B이며, 이는 Mamba가 Transformer 구성과 일치하지 않기 때문입니다.
• 모든 모델은 Mamba 논문에 설명된 Chinchilla 레시피로 훈련되며, 이는 부록 C에 재현되었습니다. Llama 아키텍처를 기반으로 한 우리의 Transformer 기준선도 Mamba 논문의 기준선을 따릅니다. 검증을 위해, 우리의 기준선은 Mamba 논문의 평가 설정에서 보고된 숫자를 재현할 수 있습니다.
• 우리의 프로토콜과 Mamba 논문의 유일한 차이점은 토크나이저입니다. Mamba 논문은 다양한 실험에서 GPT-2와 GPT-NeoX라는 두 가지 토크나이저를 사용합니다. 일관성을 위해, 우리는 이 논문 전반에 걸쳐 단일 토크나이저(Llama 토크나이저 [73])를 사용합니다.
• 우리는 하이브리드 아키텍처(예: Griffin [18])를 실험하지 않습니다. 우리의 기준선이 하이브리드가 아니기 때문입니다. 자기-주의와 TTT 레이어를 모두 사용하는 하이브리드 아키텍처는 성능을 향상시킬 수 있지만, 학문적 평가의 명확성을 떨어뜨릴 수 있습니다.

그림 11

Pile 데이터셋에서 2k 및 8k 문맥 길이에 대한 평가입니다. 자세한 내용은 3.1절에 있습니다. TTT-Linear는 2k 문맥에서 Mamba와 비슷한 성능을 보이며, 8k 문맥에서는 더 나은 성능을 보입니다.

 

그림 11로부터의 관찰

• 2k 문맥에서: TTT-Linear (M), Mamba, 그리고 Transformer는 성능이 비슷하며, 선들이 대부분 겹칩니다. TTT-MLP (M)는 큰 FLOP 예산에서는 약간 더 나쁜 성능을 보입니다. 비록 TTT-MLP가 모든 모델 크기에서 TTT-Linear보다 더 나은 퍼플렉시티를 보이지만, 추가적인 FLOP 비용이 그 이점을 상쇄합니다.
• 8k 문맥에서: TTT-Linear (M)와 TTT-MLP (M) 모두 2k 문맥에서의 관찰과는 달리 Mamba보다 훨씬 더 좋은 성능을 보입니다. Transformer 백본을 사용하는 TTT-MLP (T)도 1.3B 모델 크기에서 Mamba보다 약간 더 나은 성능을 보입니다. 이 논문 전반에 걸쳐 관찰된 강력한 현상은 문맥 길이가 길어질수록 TTT 레이어의 Mamba에 대한 이점이 더 커진다는 것입니다.
• 8k 문맥에서: Transformer는 모든 모델 크기에서 여전히 좋은 퍼플렉시티(최고는 아니더라도)를 보이지만, FLOP 비용 때문에 경쟁력이 떨어집니다.

백본의 영향

TTT 레이어를 Mamba 백본에서 Transformer 백본으로 바꾸는 것은 두 가지 효과를 가집니다. 첫째, TTT 레이어가 Mamba 백본을 사용할 때, 우리의 평가에서 더 좋은 성능을 보입니다. 둘째, Mamba 백본을 사용할 때, TTT-MLP는 최고로도 TTT-Linear와 비슷한 성능을 보이지만, Transformer 백본을 사용할 때, TTT-MLP는 명확히 더 나은 성능을 보입니다. 우리는 Mamba 백본의 시간적 컨볼루션이 시퀀스 모델링 레이어의 숨겨진 상태가 덜 표현력이 있을 때 더 큰 도움을 준다고 가정합니다. 선형 모델이 MLP보다 덜 표현력이 있기 때문에 컨볼루션의 이점을 더 많이 받습니다. 우리는 다음 절에서 이 가설을 다시 논의할 것입니다.

선형 적합의 부족

Chinchilla 논문에서는 그들의 레시피를 따르는 계산 최적 모델이 FLOP 대 퍼플렉시티의 로그-로그 플롯에서 직선에 놓인다고 경험적으로 관찰했습니다. 이는 스케일링 법칙 실험에서 흔한 경우입니다 [34]. 그러나 우리는 그림 11이나 그림 12(Books에서의 유사한 실험)에서 선형 적합을 명확히 관찰하지 못했습니다. 심지어 Transformers에서도 그렇습니다. 데이터셋, 문맥 길이, 토크나이저, 아키텍처의 차이를 고려할 때 이는 놀라운 일이 아닙니다. Mamba 논문을 따라, 큰 오차 때문에 선형 회귀로 맞추는 대신 점들을 연결합니다.

 

3.2 긴 문맥: Books

그림 12: Books에서의 문맥 길이 2k와 32k에 대한 평가

• 세부 사항은 3.2절에 있습니다. 문맥 길이 1k, 2k, 4k, 8k, 16k, 32k에 대한 전체 결과는 부록의 그림 18에 포함되어 있습니다. 대부분의 관찰은 Pile에서와 동일하게 유지됩니다.

긴 문맥에서의 성능을 평가하기 위해, 우리는 Pile의 인기 있는 하위 집합인 Books3를 사용하여 1k에서 32k까지의 문맥 길이에서 실험을 수행합니다. 여기서의 훈련 레시피는 Pile과 동일하며, TTT 레이어에 대한 모든 실험은 한 번의 훈련 실행에서 수행됩니다.

 

12 Mamba 논문을 따라, 우리는 문맥 길이에 상관없이 각 훈련 배치당 0.5M 토큰을 항상 사용합니다. 이는 문맥 길이 TTT에 대해 배치당 0.5M/T​ 시퀀스를 갖는다는 의미입니다(나눌 수 있다고 가정).

그림 12의 결과에서 몇 가지 관찰을 합니다:

• 2k 문맥에서: Books에서의 2k 문맥에서는 Pile 2k에서의 모든 관찰이 여전히 유효하지만, Mamba가 이제 TTT-Linear보다 약간 더 나은 성능을 보입니다(Pile 2k에서는 두 선이 대체로 겹쳤습니다).
• 32k 문맥에서: TTT-Linear (M)와 TTT-MLP (M) 모두 Mamba보다 더 나은 성능을 보입니다. 이는 Pile 8k에서의 관찰과 유사합니다. Transformer 백본을 사용하는 TTT-MLP (T)도 32k 문맥에서 Mamba보다 약간 더 나은 성능을 보입니다.
• 1.3B 스케일에서: TTT-MLP (T)는 TTT-MLP (M)보다 약간 더 나쁜 성능을 보입니다. 깨끗한 선형 적합의 부족으로 인해 경험적 스케일링 법칙을 도출하기는 어렵지만, TTT-MLP (T)의 강력한 추세는 Transformer 백본이 더 큰 모델과 더 긴 문맥에서 더 적합할 수 있음을 시사합니다.

우리는 LLM 훈련 비용 때문에 2k와 32k에 대해서만 백본을 소거 실험했습니다. 향후 연구에서는 더욱 표현력 있는 숨겨진 상태를 가진 TTT 레이어를 사용하면 시간적 컨볼루션이 있는 Mamba 백본이 불필요해질 것으로 생각됩니다.

Transformer 파인튜닝

Mamba 논문을 따라 Transformer를 처음부터 훈련했지만, 실제로는 긴 문맥에서는 이 접근 방식이 거의 사용되지 않습니다. 표준 방식은 짧은 문맥에서 Transformer를 훈련한 다음 긴 문맥에서 파인튜닝하는 것입니다. 이 관행을 반영하기 위해, 우리는 4k 이상의 문맥 길이에 대해 TF 파인튜닝이라는 기준선을 추가합니다. 이 기준선은 Books 2k에서 Chinchilla 레시피에 따라 훈련된 모델에서 시작하여, Llama Long 논문 [78]을 따라 지정된 문맥 길이에서 20% 더 많은 토큰을 사용하여 파인튜닝합니다. TF 파인튜닝 레시피의 자세한 내용은 부록 C에 있습니다.

문맥 길이 1k, 2k, 4k, 8k, 16k, 32k에 대한 전체 결과는 TF 파인튜닝을 포함하여 부록의 그림 18에 있습니다.

 

그림 13: 문맥 길이 1k에서 32k까지의 Books 실험

문맥 길이를 하이퍼파라미터로 취급하고 선택된 지점을 연결합니다. 우리는 처음부터 훈련된 Transformer와 파인튜닝된 Transformer를 포함하여 TF 프리트레인(TF pretrain) 및 TF 파인튜닝(TF finetune)으로 표시합니다. 그림 2의 왼쪽 패널은 350M에서 1.3B 파라미터 사이의 확대된 뷰로, 여기서 Transformer는 더 강력한 Transformer 기준선인 TF 파인튜닝을 나타냅니다.

 

문맥 길이 하이퍼파라미터로서의 역할

입력 시퀀스의 길이는 사용자가 결정하지만, 언어 모델이 입력을 처리하는 문맥 길이는 엔지니어의 설계 선택에 따라 결정됩니다. 따라서 문맥 길이는 다른 하이퍼파라미터와 마찬가지로 선택될 수 있는 하이퍼파라미터입니다.

 

정확히 말하면, 두 가지 하이퍼파라미터가 있습니다: LLM이 훈련되는 문맥 길이와 평가되는 문맥 길이입니다. 이 두 가지는 사용자가 결정하는 시퀀스 길이와 다를 수 있습니다. Transformers는 평가 문맥이 훈련 문맥보다 길 때 성능이 저하되는 경향이 있습니다. 따라서 우리는 항상 훈련 문맥 길이에서 평가를 수행하여 두 하이퍼파라미터를 동일하게 만듭니다.

선형 복잡도를 가진 LLM의 경우, 모든 문맥 길이가 동일한 FLOP를 가지므로 퍼플렉시티가 가장 낮은 값을 선택합니다. Transformers의 경우, 긴 문맥은 더 많은 FLOP를 요구하므로, 로그-로그 플롯에서 모든 점의 볼록 껍질(convex hull)을 형성하고 경계에 있는 점들을 연결합니다.

그림 13의 관찰 사항

• TTT-Linear와 TTT-MLP, 가장 성능이 좋은 방법들의 선은 거의 완전히 겹칩니다. Mamba와 TF 파인튜닝의 선도 10^20 FLOP 이후에 대부분 겹칩니다.
• TF 파인튜닝은 TF 프리트레인보다 훨씬 더 좋은 성능을 보이며, 이는 긴 문맥의 이점을 누리면서도 훈련 FLOP 비용이 매우 크지 않기 때문입니다. TF 파인튜닝과 프리트레인의 추론 FLOP는 동일하게 낮은 성능을 보이는데, 이는 이 플롯에 반영되지 않았습니다.
• 처음부터 훈련된 모든 방법(TF 프리트레인 포함)에서는 문맥 길이가 너무 길어지면 퍼플렉시티가 더 나빠집니다. 이 추세는 부록의 그림 19에서 강조됩니다. 이 추세에 대한 추가 조사는 미래 연구로 남겨둡니다.

그림 2의 왼쪽 패널은 그림 13의 확대된 뷰입니다. 명확성을 위해, 그림 2에서는 TF 프리트레인을 생략하고 더 강력한 기준선인 TF 파인튜닝(Transformer로 라벨링됨)만을 보여줍니다. 그림 14는 TTT-MLP와 추가 논의와 함께 그림 2의 오른쪽 패널을 재현한 것입니다.

 

그림 14: 그림 2의 오른쪽 패널 재현, 이제 TTT-MLP 포함

TTT-Linear와 비교할 때, TTT-MLP는 짧은 문맥에서는 약간 더 나쁜 성능을 보이지만, 긴 문맥에서는 더 나은 성능을 보입니다. 이 관찰은 숨겨진 상태로서의 MLP가 선형 모델보다 더 표현력이 있다는 우리의 예상과 일치합니다. 다시 말하지만, 모든 방법은 Mamba 1.4B와 동일한 훈련 FLOP를 가집니다. TTT-Linear와 TTT-MLP의 경우, 이 프로토콜은 동일한 추론 FLOP를 의미합니다. Transformer (TF 파인튜닝)은 2.8 × 의 추론 FLOP를 가지고 있어, 기준선으로서 이점을 제공합니다.

 

3.3 실제 시간 (Wall-clock time)

그림 15: NVIDIA A100 GPU 벤치마크

NVIDIA A100 GPU (80G HBM 및 PCIe 연결)를 사용한 벤치마크입니다.

왼쪽:

• 그림 3과 동일하지만 이제 TTT-MLP가 포함되어 있습니다. 우리의 Transformer 기준선은 Mamba 논문보다 훨씬 빠릅니다. 이는 HuggingFace Transformer [77] 대신 최첨단 서빙 시스템인 vLLM [46]을 사용했기 때문입니다.
• TTT-Linear와 Mamba는 생성(decode) 작업에서 거의 동일한 지연 시간을 가지며, 이는 Transformer와 TTT-MLP보다 훨씬 작습니다.

LLM 훈련과 추론은 순방향(forward), 역방향(backward), 생성(generate)으로 나눌 수 있습니다. 추론 중 프롬프트 처리는 전방향과 동일한 작업이며, 중간 활성화는 역방향을 위해 저장할 필요가 없습니다. 전방향(훈련 및 추론 중)과 역방향은 병렬화할 수 있으므로 우리는 이중 형식을 사용합니다. 새 토큰을 생성하는 것은 순차적이므로 원시 형식을 사용합니다.

 

리소스 제약으로 인해, 우리의 실험은 JAX로 작성되어 TPU에서 실행됩니다. v5e-256 TPU pod에서, Transformer 기준선은 문맥 2k에서 훈련당 0.30초가 걸리며, TTT-Linear는 0.27초가 걸려 시스템 최적화 없이도 이미 10% 빠릅니다. 그러나 Mamba는 GPU에서만 실행할 수 있으므로 공정한 비교를 위해 우리의 방법도 GPU에서 실행되도록 초기 시스템 최적화로 다시 작성합니다.

 

구체적으로, 우리는 ThunderKittens [66]에서 전방향을 위한 GPU 커널을 작성합니다. 역사적으로 RNN은 병렬 처리와 행렬 곱셈의 비효율적인 사용으로 인해 전방향과 역방향에서 비효율적이었습니다. 우리의 전방향 커널의 목표는 이러한 문제를 해결하기 위해 미니배치 TTT와 이중 형식의 효과를 입증하는 것입니다. 역방향을 위한 커널은 전방향과 동일한 효율성을 가지지만 수동 미분을 위한 더 복잡한 논리가 필요하므로 미래 작업으로 남겨둡니다.

 

그림 15의 왼쪽 패널은 배치 크기 16에서 전방향 커널의 지연 시간을 보여줍니다. 모든 모델은 1.3B (Mamba는 1.4B)입니다. 문맥 길이가 증가함에 따라 Transformer의 토큰당 시간이 선형으로 증가하지만, 다른 방법들은 대체로 일정하게 유지됩니다.

 

네트워크의 전방향 지연 시간이 TTT-Linear, TTT-MLP, Mamba 모두 약간 증가하는 것을 관찰했으며, 각 시퀀스 모델링 레이어의 지연 시간은 일정하게 유지됩니다. 이는 θX 연산의 지연 시간은 이론적으로 일정하지만, 실제로는 TTT가 매우 커지면 GPU 제한으로 인해 약간 증가할 수 있기 때문입니다 [30]. 우리의 Transformer 기준선이 Mamba 논문보다 훨씬 빠른 이유는 HuggingFace Transformer [77] 대신 vLLM [46]을 사용했기 때문입니다.

또한, 우리는 Triton [72]에서 생성 작업을 위한 또 다른 GPU 커널을 작성하고, 배치 크기 512에서 속도를 벤치마크합니다. 실제 시간의 또 다른 인기 있는 지표는 처리량입니다. 이는 더 큰 배치 크기를 사용할 수 있는 잠재적 이점을 고려합니다.

각 방법의 처리량을 계산하기 위해, 우리는 배치 크기를 2배씩 증가시키면서 GPU 메모리가 부족할 때까지 계속합니다. 그런 다음 각 배치 크기에서 초당 토큰 수를 측정하고 가장 높은 값을 선택합니다.

완전성을 위해, 우리는 전방향과 생성 작업에 대한 처리량을 부록의 그림 20에 보고합니다. 위의 모든 관찰과 방법 간의 순서는 처리량에서도 동일하게 유지됩니다.

 

4 관련 연구

4.1 현대 RNN

Mamba는 Structured State-Space Models [27, 21, 57, 18] 중 하나입니다. 이 모델들의 숨겨진 상태는 LSTM과 유사하게 벡터로 구성됩니다. TTT-Linear나 TTT-MLP의 경우, 숨겨진 상태는 하나 또는 두 개의 행렬이므로 더 큽니다. 그림 14에서, TTT 레이어가 더 큰 숨겨진 상태를 활용하여 긴 문맥에서 더 많은 정보를 압축할 수 있음을 알 수 있으며, 이로 인해 TTT-MLP는 TTT-Linear보다, TTT-Linear는 Mamba보다 더 우수한 성능을 보입니다.

 

TTT-Linear와 유사하게, RWKV [55, 56], xLSTM [5], Gated Linear Attention (GLA) [79]도 선형 어텐션 [41]에서 유래한 행렬 숨겨진 상태를 가집니다. GLA와 같은 현대 RNN은 청크별 병렬 처리를 사용하여 하드웨어 효율성을 향상시키므로, 청크 내의 토큰은 누적 합(cumsum) 대신 행렬 곱셈(matmuls)으로 처리될 수 있습니다. 그러나 청크별 병렬 처리는 모델의 표현력을 변경하지 않으며, 모든 시간적 종속성은 여전히 누적 합과 동일합니다.

 

 

4.2 테스트 시 학습

테스트 시 학습의 아이디어는 기계 학습에서 오랜 역사를 가지고 있습니다. 이 아이디어의 초기 버전 중 하나는 지역 학습(local learning, Bottou and Vapnik [10])입니다: 각 테스트 입력에 대해 예측을 하기 전에 이웃에서 학습하는 것입니다. 이 절차는 SVM [81]부터 현대의 LLM [29]까지 다양한 모델에 효과적으로 적용되었습니다.

 

테스트 시 학습의 또 다른 초기 버전은 전이 학습(transductive learning) [22]입니다. Vladimir Vapnik [74]에 의해 언급된 전이 원리는 "필요한 답만 얻고, 더 일반적인 답은 얻지 않는 것"입니다. 전이 학습의 실용적인 구현은 테스트 데이터를 사용하여 SVM의 여백에 제약을 추가합니다 [39, 17]. 그러나 전이 학습은 여러 테스트 인스턴스가 있어야 실질적으로 효과적입니다. 반면, 많은 테스트 시 학습 인스턴스는 단일 테스트 인스턴스(이미지, 비디오 또는 자연어 시퀀스)만 필요합니다.

 

컴퓨터 비전에서는 테스트 시 학습의 아이디어가 수십 년 동안 얼굴 인식 [38], 객체 탐지 [53], 이미지 초해상도 [65], 3D 재구성 [50]과 같은 응용 프로그램에 적용되었습니다. 최근에는 자연어 처리에도 동일한 아이디어가 적용되어 동적 평가(dynamic evaluation)라고 불립니다 [44, 45]. 기본 접근 방식은 테스트 시퀀스에서 언어 모델을 직접 미세 조정하는 것입니다.

 

다음으로, 테스트 시 학습과 빠른 가중치(fast weights)에 대한 두 가지 관련 연구 라인을 자세히 논의합니다.

 

 

4.2.2 빠른 가중치

빠른 가중치(fast weights)의 일반적인 아이디어는 모든 데이터를 사용하는 대신 가장 관련성 높은 데이터에만 "빠른" 모델의 파라미터를 업데이트하는 것입니다 [71]. 이 아이디어는 1980년대부터 존재해왔습니다 [32]. 가장 관련성 높은 데이터는 테스트 인스턴스 자체일 수 있으므로, TTT는 빠른 가중치의 특수한 경우로 볼 수 있습니다.

 

빠른 가중치의 이전 연구는 일반적으로 데이터의 일부 목표를 최적화하는 명시적인 학습 문제를 형성하는 것을 피합니다. 예를 들어, Hebbian 학습과 Hopfield 네트워크의 업데이트 규칙은 각 입력 x에 대해 단순히 x x^T (또는 그 변형)을 빠른 가중치에 추가합니다 [4]. 반면에, TTT는 명시적인 학습 문제를 공식화하는 아이디어를 수용하며, 여기서 테스트 인스턴스는 일반화의 대상입니다. 우리의 업데이트 규칙도 명시적인 최적화 단계입니다.

 

빠른 가중치 프로그래머(FWPs)의 아이디어는 "느린" 모델로 빠른 가중치를 업데이트하는 것입니다 [62]. 우리의 내부 루프 가중치 W는 "빠른" 가중치로, 외부 루프 가중치 θ는 "느린" 가중치로 볼 수 있습니다. 따라서 TTT 레이어를 포함한 네트워크는 FWPs의 특수한 경우로 볼 수 있습니다 [43]. FWPs의 Hebbian 업데이트 규칙은 선형 어텐션과 동등하며 [60], 따라서 배치 기울기 하강을 사용하는 단순한 TTT-Linear와도 동등합니다.

 

FWPs의 정의는 매우 넓습니다. 실제로, SwiGLU 블록이 있는 Transformer와 같은 게이팅 메커니즘을 가진 모든 네트워크도 FWPs의 특수한 경우로 볼 수 있습니다. 최근 연구는 언어 모델링을 위해 FWPs를 실험하고 있습니다: Irie et al. [37]는 "느린" 네트워크의 출력으로 생성된 가중치를 가진 "빠른" 네트워크를 설계했습니다. Clark et al. [16]은 초기화가 느린 가중치로 학습된 빠른 가중치의 마지막 레이어를 Transformer에 추가했습니다. 기존의 FWPs 연구와 비교하여 우리의 기여는 업데이트를 위한 명시적인 학습 문제를 공식화하는 것이며, 이는 미니배치와 LN과 같은 학습 도구를 사용할 수 있게 해줍니다.

 

4.3 학습하는 방법을 배우기 (Learning to Learn)

수십 년 동안 연구자들은 학습하는 방법을 배우기, 즉 메타 학습 또는 이중 최적화가 지능의 중요한 요소가 되어야 한다고 주장해왔습니다 [61, 6, 70, 47]. 이전 연구들 [2], [20], [52]에서는 내부 루프가 한 번에 전체 데이터셋에서 학습하므로, 외부 루프는 데이터셋 또는 작업의 집합이 필요했습니다. 요약하면, 외부 루프는 일반적인 훈련보다 "한 단계 위"에 있습니다. 수백만 개의 데이터셋을 수집하는 것이 어렵기 때문에, 이러한 외부 루프는 확장하기 어렵습니다.

 

반면에, TTT의 경우 각 시퀀스 자체가 하나의 데이터셋이자 고유한 일반화 문제를 정의합니다. 내부 루프는 일반적인 훈련보다 "한 단계 아래"에 있기 때문에, 우리의 외부 루프는 데이터셋 간의 일반화와 같은 새로운 문제 설정이 아니라 지도 학습의 정형화된 문제에 대한 또 다른 해결책일 뿐입니다. 표 2에서 설명한 것처럼, 우리의 외부 루프는 일반적인 훈련과 "같은 수준"에 있습니다. 이는 우리의 외부 루프를 확장하기 쉽게 만듭니다.

 

표 2: 요약하면, 우리의 논문은 지도 학습을 두 개의 중첩된 루프가 있는 학습하는 방법을 배우기로 재구성합니다. 강조된 외부 루프의 행들은 일반적인 훈련과 동일합니다. 외부 루프의 파라미터는 내부 루프의 하이퍼파라미터가 됩니다. 직관적으로, 내부 루프(TTT)는 일반적인 훈련보다 "한 단계 아래"에 있습니다.

 

5. 토론

우리는 정형화된 지도 학습 문제를 학습하는 방법으로 (테스트 시 학습) 재구성했습니다. 우리의 공식화는 전통적으로 네트워크 아키텍처라고 알려진 것을 구축하기 위한 대체 개념 프레임워크를 제공합니다. 표 2에서 현재의 구현을 요약합니다.

이 프레임워크 내에서 효과적인 구현을 찾기 위한 탐색 공간은 거대하며, 우리의 논문은 그 중 첫걸음에 불과합니다. 다행히도 우리의 관점이 유효하다면, 정규 훈련에서 사용되는 휴리스틱을 테스트 시 학습으로 이전할 수 있으며, 탐색이 효율적일 수 있습니다. 다음은 미래 연구를 위한 특히 유망한 방향을 간략히 설명합니다.

• 외부 루프 매개변수화: 다중 뷰 재구성 작업의 가족을 매개변수화하는 다른 방법이 많이 있습니다. 또는 아마도 더 일반적인 자기 지도 작업의 가족일 수도 있습니다. 우리가 처음 시도한 것이 최선이라면 그것은 큰 우연일 것입니다.
• 시스템 최적화: 3.3절에서의 시스템 최적화는 초기 단계일 뿐이며, 이를 개선할 방법이 많이 있습니다. 또한, 시간에 따른 파이프라인 병렬 처리를 통해 여러 장치에서 수백만 개의 토큰으로 구성된 긴 시퀀스를 함께 처리할 수 있을 것입니다.
• 더 긴 문맥 및 더 큰 모델: 우리의 학문적 자원에 의해 제한되어, 수백만 또는 수십억의 문맥 길이로 훈련하지 않았습니다. 이는 그림 19에 따르면 더 큰 모델이 필요할 수 있습니다. TTT 레이어의 이점은 더 긴 문맥에서 더 두드러질 것입니다.
• f의 더 야심찬 구현: 문맥 길이가 길어짐에 따라 f도 더 커져야 할 것입니다. 비디오 작업 및 구체화된 에이전트의 경우, 문맥 길이가 쉽게 수백만 또는 수십억으로 확장될 수 있으며, f는 컨볼루션 신경망일 수 있습니다.
• 다중 레벨 학습 학습: f 자체가 자기 주의 레이어인 경우, 정리 2에 따르면 이는 기존 내부 루프에 중첩된 또 다른 내부 루프로 해석될 수 있습니다. 이러한 방식으로 우리는 잠재적으로 여러 레벨의 중첩된 학습 문제를 구축할 수 있습니다.

TTT를 연구하는 이유

먼저 더 기본적인 질문: 왜 AI를 연구합니까? 우리 중 일부에게 AI는 인간 지능의 본질을 탐구하기 위한 놀이터입니다. 이전 연구는 종종 기계 학습으로 인간 학습을 모델링하려고 합니다. 여기서 훈련은 i.i.d. 인스턴스로 섞인 데이터셋에서 이루어지며, 추론은 별도의 테스트 세트에서 이루어집니다. 그러나 인간은 자연스럽게 i.i.d. 인스턴스로 학습하거나 훈련-테스트 분할을 하지 않습니다. 우리는 인간 학습이 TTT, 즉 내부 루프와 더 유망한 연결이 있다고 믿습니다. 이 내부 루프의 데이터는 잠재적으로 매우 긴 시퀀스로 강한 시간적 종속성을 가지며, 데이터의 어느 부분이든 훈련과 테스트에 사용될 수 있습니다. 이것이 우리가 TTT를 연구하는 이유입니다.

 

저자 기여

• Yu Sun: 2022년 11월 Xinhao Li와 함께 이 프로젝트를 시작했고, 2023년 6월부터 전업으로 일해왔습니다. Yu는 프로젝트의 개념 프레임워크를 제안하고, 미니배치 TTT와 이중 형식을 설계했으며, 다른 사람들의 도움을 받아 논문을 작성하고 팀의 일상 운영을 이끌었습니다.
• Xinhao Li: 2022년 11월 Yu Sun과 함께 이 프로젝트를 시작했고, 그 이후로 전업으로 일해왔습니다. Xinhao와 Karan은 현재 코드베이스 개발을 공동으로 이끌었습니다. 2024년 3월 이전에는 Xinhao가 이 프로젝트의 방향을 형성한 초기 코드베이스의 주요 기여자였습니다.
• Karan Dalal: 2023년 6월에 전업으로 이 프로젝트에 합류했습니다. Xinhao와 협력하여 현재 코드베이스 개발을 공동으로 이끌었고, 3장에서 실험을 관리하고 논문 작성에 도움을 주었으며, 프로젝트 방향에 중요한 기여를 했습니다.
• Jiarui Xu: 2024년 3월에 이 프로젝트에 합류했습니다. 합류 이후 우리의 아키텍처 개발을 이끌었고, 프로젝트 방향에 중요한 기여를 했습니다.
• Arjun Vikram: 2023년 9월에 이 프로젝트에 합류했습니다. 시스템 최적화 및 현재와 이전 코드베이스에 중요한 기여를 했습니다.
• Genghan Zhang: 2024년 1월에 이 프로젝트에 합류했습니다. 중요한 통찰을 제공하고 시스템 최적화에 중요한 개선을 했습니다.
• Yann Dubois: 2024년 2월에 이 프로젝트에 합류했습니다. 현재의 fff 구현을 제안하고, 프로젝트 방향에 중요한 기여를 했습니다.
• Xinlei Chen와 Xiaolong Wang: 2022년 11월부터 이 프로젝트를 지원해왔으며, 오랜 기간 동안 테스트 시 학습의 방향을 지원해왔습니다. 초기 단계에서의 계산 및 조직 지원 없이는 이 프로젝트가 생존할 수 없었습니다. 실험에 귀중한 조언을 제공했습니다.
• Sanmi Koyejo, Tatsunori Hashimoto, Carlos Guestrin: 2023년 5월부터 이 프로젝트를 지원해왔습니다. 실험 및 프레젠테이션에 귀중한 조언을 제공했습니다. 예를 들어, Sanmi는 우리가 TTT-Linear에 집중하도록 제안했고, Tatsu는 그림 2의 (왼쪽) 실험을 제안했으며, Carlos는 2장을 개요했습니다.

감사의 글

이 프로젝트의 일부 계산은 Google TPU Research Cloud 프로그램의 관대한 지원을 받았습니다. XW는 Amazon Research Award, Cisco Faculty Award 및 Qualcomm Innovation Fellowship의 일부 지원을 받았습니다. SK는 NSF 2046795 및 2205329, NIFA award 2020-67021-32799, Alfred P. Sloan Foundation 및 Google Inc.의 지원을 받습니다. TH는 Sony Faculty Innovation Award와 Panasonic의 기부금을 받습니다. CG는 Air Force Office of Scientific Research (AFOSR), FA9550-20-1-0427, Stanford Human-Centered Artificial Intelligence (HAI) Institute, Google 및 IBM의 기부금의 지원을 받습니다.

 

우리는 Rohan Taori, Xuechen Li, Allan Zhou, Ke Chen 및 Guandao Yang과 많은 유익한 논의를 해주신 분들께 감사드립니다. Menghao Guo에게 코드 릴리스 도움, Xinyang Geng에게 EasyLM 도움, Hao Liu에게 LWM 코드베이스 도움, David Hall에게 Levanter 도움, Yossi Gandelsman 및 Yutong Bai에게 프로젝트 초기 단계 도움, Mert Yuksekgonul에게 논문 도표 도움, Horace He 및 Azalia Mirhoseini에게 시스템 도움, Sharad Vikram 및 Roy Frostig에게 JAX 및 Pallas에 관한 질문에 대한 답변, Albert Gu 및 Tri Dao에게 Mamba 논문의 실험 재현 도움, Kilian Weinberger 및 Percy Liang에게 프레젠테이션 조언을 주신 것에 감사드립니다. Yu Sun은 PhD 지도 교수인 Alexei A. Efros와 Moritz Hardt에게 이 논문의 일부가 된 여러 해전의 통찰에 대해 감사드립니다.

 

2407.04620v1.pdf

1.63MB