[ML] Titans: Learning to Memorize at Test Time by Google

소개

2017년, Google은 오늘날 우리가 경험하는 AI 혁명을 촉발한 획기적인 논문인 “Attention is All You Need“를 발표했습니다. 이 논문에서는 오늘날 전부는 아니더라도 대부분의 상위 대형 언어 모델(LLM)의 근간이 된 트랜스포머를 소개했습니다.

트랜스포머(Transformer)의 강점과 비용

트랜스포머의 강점은 주로 attention 사용에서 기인합니다. 토큰 시퀀스가 주어지면 Transformers는 전체 시퀀스를 한 번에 처리하고, attention 메커니즘을 사용해 전체 시퀀스의 종속성을 파악하여 고품질의 결과물을 제공합니다. 하지만 이 강력한 기능에는 입력 시퀀스 길이에 대한 이차 종속성이라는 비용이 따릅니다. 이 비용은 트랜스포머가 더 긴 시퀀스로 확장하는 데 제한을 초래합니다.

반복 모델 (Recurrent Models)

반면에 반복 모델은 동일한 이차 종속성을 겪지 않습니다. 전체 시퀀스를 한 번에 처리하는 대신, 시퀀스의 데이터를 숨겨진 상태라고도 하는 압축 메모리로 압축하여 점진적으로 처리합니다. 이러한 선형 종속성은 반복 모델의 확장성을 향상시키는 데 기여합니다. 그러나 반복 모델은 트랜스포머만큼 성능이 뛰어난 것으로 입증되지는 않았습니다.

타이탄즈(Titans) 소개

이 게시물에서는 “Titans: Learning to Memorize at Test Time” 이라는 새로운 모델 아키텍처를 소개하는 Google Research의 새 논문을 살펴봅니다. 타이탄즈이라는 새로운 모델 아키텍처를 소개하는 이 논문에서는 트랜스포머의 이차적 비용 문제를 완화하면서 유망한 결과를 보여주는 타이탄즈이라는 새로운 모델 아키텍처를 소개합니다. 타이탄즈 모델은 인간의 뇌에서 기억이 작동하는 방식에서 영감을 받아 설계되었습니다. 논문에서 흥미로운 인용문을 보면 기억은 기본적인 정신 과정이며 인간 학습에서 분리할 수 없는 요소라고 언급하고 있습니다. 기억 시스템이 제대로 작동하지 않는다면 인간과 동물은 기본적인 반사 신경과 정형화된 행동에만 국한될 것입니다.

타이탄즈 논문 작성자

심층 신경망 장기기억 모듈 (Deep Neural Long-Term Memory Module)

타이탄즈 논문의 핵심적인 기여는 심층 신경망 장기기억 모듈입니다. 먼저 심층 신경망 장기기억 모듈이 무엇인지 이해한 다음, 이 모듈이 타이탄즈 모델에 어떻게 통합되는지 살펴보겠습니다.

메모리가 고정된 벡터로 인코딩되는 순환 신경망과 달리 심층 신경 장기기억 모듈은 과거 기록의 추상화를 매개변수로 인코딩하는 모델, 즉 여러 계층으로 구성된 신경망입니다. 이러한 모델을 훈련하기 위해 모델이 훈련 데이터를 암기하도록 훈련하는 것이 한 가지 아이디어입니다. 그러나 암기는 모델의 일반화를 제한하고 성능이 저하될 수 있는 것으로 알려져 있습니다.

과적합 없는 암기 (Memorization Without Overfitting)

연구진은 암기할 수 있는 모델을 만들되, 학습 데이터에 모델을 과적합하지 않는 흥미로운 접근 방식을 고안했습니다. 이 접근법은 인간 기억의 비유에서 영감을 얻었습니다. 우리는 우리를 놀라게 하는 사건을 접할 때 그 사건을 기억할 가능성이 더 높습니다. 신경 장기기억 모듈의 학습 과정은 이를 반영하도록 설계되었습니다.

서프라이즈(surprise) 모델링

서프라이즈(surprise) 구성 요소를 사용하여 신경 모델 업데이트하기

연구진이 위의 정의를 사용하여 서프라이즈을 어떻게 모델링했는지 논문을 통해 알아볼 수 있습니다. $M_t$는 $t$ 시점의 신경 장기기억 모듈을 나타내며, 이전 시간 단계의 파라미터와 기울기로 모델링된 서프라이즈 요소를 사용하여 업데이트됩니다. 기울기가 크면 모델은 입력에 더 많이 놀라게 되고, 그 결과 모델 가중치가 더 크게 업데이트됩니다.

그러나 이 정의는 모델이 서프라이즈 순간이 발생한 직후에 발생하는 중요한 정보를 놓칠 수 있으므로 여전히 이상적이지 않습니다.

과거의 놀라움(past surprise) 모델링

인간의 관점에서 볼 때, 놀라운 사건은 기억에 남지만 오랜 시간이 지나도 계속해서 우리를 놀라게 하지는 않습니다. 우리는 보통 놀라운 사건에 적응합니다. 그럼에도 불구하고 그 사건은 오랜 시간 동안 우리의 주의를 끌 정도로 충분히 놀라웠을 수 있으며, 이로 인해 전체 기간을 기억하게 됩니다.

과거(past) 및 순간적(momentary) 놀라움(suprise)을 모두 사용하여 신경 기억 업데이트하기

위의 정의에서 과거 놀라움의 모델링을 포함하는 개선된 모델링에 대해 배울 수 있습니다. 이제 이전 가중치의 상태를 사용하여 신경 장기기억의 가중치를 업데이트하고 $S_t$로 표시된 서프라이즈 구성 요소를 업데이트합니다. 서프라이즈 구성 요소는 이제 시간에 따라 측정되며 이전 서프라이즈과 이전 섹션에서 이미 설명한 것과 동일한 순간적 서프라이즈과 감쇠 계수로 구성됩니다.

아직 모델링되지 않은 또 다른 요소는 망각(Forgetting)입니다.

망각(forgetting) 모델링하기

모델이 과거 정보를 잊도록 허용하기

예를 들어 수백만 개의 토큰과 같은 매우 큰 시퀀스를 다룰 때는 어떤 과거 정보를 잊어야 하는지 관리하는 것이 중요합니다. 위의 정의에서 최종 모델링을 확인할 수 있습니다. 이러한 정의는 이전 섹션의 정의와 동일하지만 알파벳으로 표시된 적응형 망각(adaptive forgetting) 메커니즘을 추가하고 게이팅(gating) 메커니즘이라고도 합니다. 이를 통해 메모리는 더 이상 필요하지 않은 정보를 잊어버릴 수 있습니다.

손실 함수(the loss function)

신경 기억 손실 함수

손실 함수는 위의 방정식으로 정의됩니다. 손실은 과거 데이터를 키와 값의 쌍으로 저장하여 연관 기억을 모델링하고 모델에 키와 값 사이의 매핑을 학습시키는 것을 목표로 합니다. 트랜스포머와 마찬가지로 선형 레이어는 입력을 키와 값으로 투영합니다. 그런 다음 손실은 메모리 모듈이 키와 값 사이의 연관성을 얼마나 잘 학습하는지를 측정합니다.

명확하게 설명하자면, 모델은 전체 시퀀스를 한 번에 처리하지 않고 점진적으로 처리하여 메모리 정보를 가중치에 축적합니다.

Titans Architecture #1 - 컨텍스트로서의 메모리(Memory as a Context, MAC)

이 논문에서는 몇 가지 다른 아키텍처를 제안합니다. 첫 번째는 컨텍스트로서의 메모리(Memory as a Context, MAC)입니다. 논문의 다음 그림을 통해 이에 대해 알아볼 수 있습니다.

Titans Memory as a Context Architecture

이 모델은 영구 메모리, 코어, 컨텍스트 메모리의 세 가지 구성 요소로 이루어져 있습니다. 긴 입력 시퀀스가 주어지면 이를 더 작은 시퀀스로 분할합니다. 코어 컴포넌트에 표시된 시퀀스는 실제로 청크된 시퀀스입니다.

영구 메모리(persistent memory)

먼저 시퀀스의 시작 부분에 학습 가능한 데이터 독립적 가중치를 추가하여 영구 메모리를 통합합니다. 이는 실제로는 입력에 영향을 받지 않는 전역 정보를 저장하는 학습 가능한 토큰입니다. 또한 시퀀스는 attention block을 통해 전달되며, 시퀀스의 시작 부분에 학습 가능한 토큰을 추가하면 편향으로 인해 초기 토큰에 주의 가중치가 매우 활성화되는 attention 누수를 방지하는 데 도움이 됩니다. 학습 가능한 메타 토큰을 추가하는 것은 엔비디아의 Hymba에서도 성공적으로 입증되었습니다.

컨텍스트 메모리(contextual memory)

컨텍스트 메모리는 컨텍스트에 따라 달라지기 때문에 심층 신경 장기기억 모듈로 구성됩니다. 입력 청크 시퀀스를 입력으로 하여 신경 장기기억 모듈의 순방향 패스를 사용하여 장기기억에서 메모리 토큰을 가져옵니다. 신경망 메모리는 테스트 시간에도 계속 업데이트되므로, 긴 시퀀스의 청크를 처리할 때 신경망 메모리는 이미 동일한 시퀀스의 이전 청크의 정보로 업데이트되어 있습니다. 검색된 메모리 토큰은 영구 메모리 토큰 다음에 입력 시퀀스에 추가됩니다.

핵심 구성 요소(the core component)

핵심 컴포넌트는 영구 메모리, 컨텍스트 메모리 및 입력 자체의 정보를 포함하는 확장된 시퀀스로 공급되는 attention block을 사용하여 이 모든 것을 하나로 통합합니다. 이렇게 하면 모델이 여러 메모리 유형을 활용할 수 있습니다. attention block의 출력은 neural memory를 업데이트하는 데 사용됩니다. 주의 메커니즘은 주어진 입력에 대해 장기기억을 사용할지 여부를 결정하는 데 도움이 됩니다. 또한 attention은 장기기억이 현재 맥락에서 유용한 정보만 저장하도록 돕습니다. 최종 출력은 attention block 출력과 신경 기억의 출력에 따라 결정됩니다.

Titans Architecture #2 - 메모리로서의 게이트(Memory as a Gate, MAG)

다음 타이탄즈 아키텍처 버전은 MAG(Memory as a Gate)라고 하며, 아래에서 볼 수 있는 논문의 멋진 그림이 있습니다.

Titans Memory as a Gate Architecture

이 버전에서는 영구 메모리, 코어, 컨텍스트 메모리를 나타내는 3개의 브랜치도 있습니다. 이전 버전과 다른 점은 시퀀스가 청크화되지 않았다는 점입니다. 여기서 입력 시퀀스는 전체 입력 시퀀스입니다. 이는 attention block에서 sliding window attention 기능을 활용함으로써 가능합니다. 영구 메모리 학습 가능한 가중치가 시퀀스의 시작 부분에 다시 추가됩니다. 하지만 이전과 달리 neural memory는 attention block의 컨텍스트에 데이터를 제공하지 않습니다. 대신, neural memory는 입력 시퀀스에서 업데이트되고 그 출력은 게이팅 메커니즘을 사용하여 핵심 분기와 결합됩니다.

Titans Architecture #3 - 계층형 메모리(Memory as a Layer, MAL)

타이탄즈 아키텍처의 세 번째 변형은 계층형 메모리(Memory as a Layer, MAL)라고 합니다. 다음 그림에서 이에 대해 알아볼 수 있습니다.

Titans Memory as a Layer Architecture

이전 버전과 유사하게 시퀀스는 청크가 아니며 sliding window attention을 사용합니다.

이 버전에서는 neural memory를 모델 레이어로 사용하여 입력 시퀀스가 학습 가능한 가중치와 함께 먼저 neural memory를 통과한 후 attention block을 통과합니다.

이 설계는 일반적으로 트랜스포머 레이어를 쌓는 방식과 유사하게 여러 개의 neural memory 모듈과 attention block으로 구성된 레이어를 쌓을 수 있습니다. 그러나 순차적 설계는 각 레이어의 파워에 따라 모델의 파워가 제한됩니다.

이 역시 같은 이유로 서로 다른 구성 요소를 순차적이지 않고 병렬로 배치한 Hymba 논문과 유사한 관찰 결과입니다.

Titans Architecture #4 - LMM

마지막 타이탄즈의 변형은 LMM으로, 이전 버전과 동일하지만 attention block이 없고 메모리 모듈에만 의존하는 방식입니다.

Results

Language Tasks Comparison

Comparison of Titans with baselines on language tasks

위의 표에서는 언어 모델링 및 상식 추론 과제에서 기준선과 타이탄즈 모델을 비교한 결과를 볼 수 있습니다. 최상의 결과는 recurrent 및 attention 메커니즘을 모두 활용하는 하이브리드 모델의 경우 파란색으로 표시되어 있고, 비하이브리드 모델의 경우 노란색으로 표시되어 있습니다.

Titan LMM은 다른 비하이브리드 모델에 비해 최고의 결과를 달성하여 신경 장기 기억 모듈의 성능을 보여줍니다. 하이브리드 모델 중에서는 MAC Titan이 전반적으로 가장 좋은 결과를 얻었으며, Memory as a Gate Titan이 약간 뒤처졌습니다.

Needle in a Haystack

Needle in a Haystack Comparison

또 다른 흥미로운 비교는 건초 더미에서 바늘 찾기 작업입니다. 위의 표에서 논문의 결과를 확인할 수 있습니다. 이 과제에서 모델은 매우 긴 텍스트에서 정보를 검색해야 하므로 모델의 실제 유효 문맥 길이를 측정해야 합니다.

제목의 숫자는 평가된 시퀀스의 길이를 나타냅니다. 세 가지 벤치마크 모두에서 시퀀스 길이가 증가함에 따라 기준선과 비교했을 때 타이탄즈의 확실한 우위를 확인할 수 있습니다.

BABILong Benchmark

Comparison on the BABILong benchmark of Titans MAC with top models

또 다른 흥미로운 결과는 위의 그림에서 볼 수 있는데, 이는 BABILong 벤치마크에서 타이탄즈과 상위 모델을 비교한 것입니다. 이 벤치마크는 매우 긴 문서에 분산된 사실을 추론해야 하는 긴 시퀀스에 대해 더 어려운 벤치마크입니다.

X축에는 시퀀스 길이가 표시되고 Y축에는 각 모델의 정확도가 측정됩니다. 빨간색 선으로 표시된 MAC Titan의 결과는 매우 긴 시퀀스에서 다른 모델보다 훨씬 뛰어난 성능을 보여줍니다!