Published on

구글 타이탄 아키텍처 트랜스포머 메모리 병목 현상 해결

저자
  • avatar
    이름
    Ajax
    Twitter

구글의 새로운 아키텍처 Titan 소개

기술 업계는 구글에서 새롭게 등장한 아키텍처 Titan에 대한 관심이 뜨겁습니다. Titan은 특히 메모리 처리 방식에서 트랜스포머 모델의 한계에 도전하도록 설계되었습니다. 구글 내부 팀에서 개발했다는 점 때문에 트랜스포머의 잠재적인 후계자로 주목받고 있습니다.

기존 모델의 메모리 문제점

LSTM 및 트랜스포머와 같은 기존 모델은 혁신적이지만 인간과 유사한 메모리를 시뮬레이션하는 데 어려움을 겪고 있습니다. 이러한 어려움은 다음과 같습니다.

  • 제한된 용량: 데이터는 고정 크기의 숨겨진 상태로 압축되어 유지할 수 있는 정보량이 제한됩니다.
  • 높은 계산 비용: 장거리 종속성을 포착할 수 있지만, 시퀀스 길이가 길어질수록 계산 비용이 제곱으로 증가하여 매우 긴 시퀀스에 비효율적입니다.
  • 훈련 데이터에 대한 과도한 의존: 단순히 훈련 데이터를 암기하는 것만으로는 테스트 데이터가 훈련 분포에서 벗어날 수 있는 실제 응용 프로그램에 항상 도움이 되지 않습니다.

Titan의 접근 방식: 신경 영감 메모리 모듈

Titan 팀은 신경망 매개변수에 정보를 인코딩하는 다른 접근 방식을 취했습니다. 그들은 테스트 중에 특정 데이터를 기억하고 잊는 방법을 배우도록 설계된 온라인 메타 모델을 개발했습니다. 이 모델은 신경 심리학적 원리에 영감을 받아 다음과 같은 주요 요소를 통합합니다.

  • 놀라움(Surprise)을 트리거로: 예상치 못한 사건은 더 쉽게 기억됩니다. "놀라움"은 메모리 모듈에 대한 입력의 기울기로 측정됩니다. 기울기가 클수록 입력이 더 예상치 못한 것입니다.
  • 모멘텀 및 망각 메커니즘: 모멘텀 메커니즘은 단기적인 놀라움을 장기 기억으로 축적하는 반면, 망각 메커니즘은 오래된 기억을 지워 메모리 오버플로를 방지합니다.
  • MLP 기반 메모리: 메모리 모듈은 여러 MLP 레이어로 구성되어 데이터의 심층 추상화를 저장할 수 있어 기존의 매트릭스 기반 메모리보다 강력합니다.

이러한 온라인 메타 학습 접근 방식은 모델이 단순히 훈련 데이터를 암기하는 것보다 새로운 데이터에 적응하는 방법을 배우는 데 집중하도록 도와줍니다. 또한 이 모듈은 병렬 계산을 위해 설계되어 효율성을 향상시킵니다.

딥러닝 아키텍처에 메모리 모듈 통합

Titan 연구팀은 딥러닝 아키텍처에 메모리 모듈을 통합하기 위한 세 가지 변형을 제안했습니다.

  1. MAC (Memory as Context): 이 방법은 장기 및 지속적 메모리(작업 지식을 인코딩)를 어텐션 메커니즘에 입력되는 컨텍스트로 결합합니다.
  2. MAG (Memory as Gate): 이 접근 방식은 두 분기에서 슬라이딩 윈도우 어텐션 메커니즘과 메모리 모듈의 게이트 융합을 사용합니다.
  3. MAL (Memory as Layer): 여기서 메모리 모듈은 어텐션 메커니즘에 공급하기 전에 과거 정보를 압축하는 독립적인 레이어로 구현됩니다.

연구팀은 각 변형에 강점과 약점이 있다는 것을 발견했습니다.

Titan의 성능 및 장점

Titan은 언어 모델링, 상식 추론, 시계열 예측을 포함한 다양한 작업에서 뛰어난 성능을 입증했습니다. 트랜스포머 및 Mamba와 같은 최첨단 모델을 능가했습니다. 특히 장기 메모리 모듈(LMM)만으로도 여러 작업에서 기준 모델보다 뛰어난 성능을 보여주어 단기 메모리(어텐션) 없이도 독립적인 학습 능력을 입증했습니다.

긴 텍스트에서 세밀한 단서를 찾는 "건초 더미에서 바늘 찾기" 테스트에서 Titan은 시퀀스 길이가 2k에서 16k로 증가하더라도 약 90%의 정확도를 유지했습니다. 연구팀은 표준 테스트가 긴 텍스트 처리에서 Titan의 장점을 완전히 보여주지 못한다고 지적합니다. 또한 Titan은 매우 긴 문서에 흩어져 있는 사실에서 추론해야 하는 작업에서 GPT4, Mamba, 심지어 RAG를 사용한 Llama3.1보다 뛰어난 성능을 보였습니다.

Titan은 시계열 예측 및 DNA 시퀀스 모델링과 같은 특정 영역에서도 인상적인 성능을 보여주었습니다.

Titan을 개발한 팀

이 연구는 현재 Google DeepMind에 속하지 않은 Google Research NYC 알고리즘 및 최적화 그룹의 팀에서 수행했습니다.

  • Ali Behrouz는 Cornell University 인턴으로 논문의 첫 번째 저자입니다.
  • Zhong Peilin은 칭화대 졸업생이자 컬럼비아 대학교 박사 졸업생으로 2021년부터 Google 연구 과학자로 재직 중입니다. 그는 학부생 시절인 STOC 2016에서 첫 번째 저자 논문을 발표한 것으로 유명합니다.
  • Vahab Mirrokni는 Google Fellow이자 부사장으로 팀을 이끌고 있습니다.

연구팀은 Pytorch와 Jax를 사용하여 Titan을 개발했으며, 곧 훈련 및 평가 코드를 공개할 계획입니다.