mini-AGI: 소비자 하드웨어를 위한 지속 학습 바이트 수준 모델

mini-AGI는 단일 소비자용 GPU(8GB VRAM)에서 처음부터 훈련할 수 있도록 설계된 바이트 수준 언어 모델입니다. 전통적인 고정 모델과 달리, mini-AGI는 온라인 학습에서 흔히 발생하는 치명적인 망각(catastrophic forgetting) 없이 데이터 스트림으로부터 지속적으로 학습할 수 있는 지속 학습 시스템을 구현합니다.

아키텍처: 페이지 기반 전문가 혼합(MoE) 및 적응형 깊이

mini-AGI는 메모리 효율성과 동적 계산을 우선시하는 전통적인 트랜스포머 아키텍처와는 다른 구조를 사용합니다. 고정된 레이어 스택 대신, 두 개의 밀집 프리루드 블록과 최대 24회 반복 적용되는 순환 블록을 사용합니다.

동적 계산 및 라우팅

  • 적응형 깊이(PonderNet): 모델은 각 문자에 대해 얼마나 많은 계산을 할지 결정하기 위해 정지 헤드(halting head)를 사용합니다. 간단한 문자는 단일 계산 행만 필요할 수 있지만, 복잡한 문자는 최대 24회까지 걸릴 수 있습니다. 이는 입력의 난이도에 따라 계산 자원을 할당함으로써 효율성을 보장합니다.
  • 순환 MoE 라우팅: 순환 블록이 적용될 때마다 모델은 공유 풀에서 상위 8개의 전문가를 선택합니다. 라우팅이 문자 단위가 아니라 블록 적용 단위이기 때문에, 한 문자는 다양한 깊이에서 다양한 전문가를 활성화할 수 있습니다.
  • 바이트 수준 처리: 모델은 256개의 바이트 값에 직접 작동합니다. 토크나이저를 제거함으로써, 새로운 어휘 없이도 어떤 데이터 유형도 읽을 수 있습니다.

페이지 기반 메모리 관리

가용 VRAM보다 큰 모델을 맞추기 위해 mini-AGI는 디스크 공간을 가중치의 주 저장소로 사용하는 페이지 시스템을 구현합니다:

  • 디스크 저장: 모든 전문가의 가중치와 Adam 최적화기 모멘텀은 디스크에 개별 파일로 저장됩니다.
  • VRAM 작업 세트: 언제나 VRAM에 소수의 전문가(작업 세트)만 존재합니다.
  • 요청 기반 스왑: 모델은 이전 청크의 라우팅 패턴을 기반으로 다음 텍스트 청크에 필요한 전문가를 예측합니다. 전문가들은 디스크, RAM 캐시, VRAM 간에 스왑되어 전방 전파를 위한 필요한 파라미터가 항상 사용 가능하도록 보장합니다.

치명적인 망각 방지

단일 스트림의 데이터로부터의 지속 학습은 종종 "치명적인 망각"을 초래합니다. 즉, 새로운 정보가 오래된 지식을 덮어쓰는 현상입니다. mini-AGI는 두 가지 주요 메커니즘을 통해 이를 해결합니다.

트렁크 학습률 분리

모델의 "트렁크"(임베딩, 어텐션, 라우터, 정지 헤드 포함)는 전문가보다 훨씬 낮은 학습률(0.1x)로 훈련됩니다.

실험 데이터에 따르면, 트렁크와 전문가에 동일한 학습률을 사용할 경우 큰 손실(망각)이 발생하지만, 0.1x 트렁크 학습률은 우연에 대한 진전의 99.84%를 유지합니다. 이는 핵심 라우팅 및 구조적 논리가 단일 데이터 스트림의 특정 콘텐츠에 의해 덮어쓰이는 것을 방지합니다.

구조적 격리

모델이 전문가 혼합(MoE) 방식을 사용하기 때문에, 어떤 훈련 단계에서도 모델의 일부만 업데이트됩니다. 524,000자 분량의 체스 데이터를 탐색하는 실험에서, 136개의 전문가 중 단지 54개만 기울기(gradient)를 받았습니다. 이 구조적 격리는 새로운 주제를 학습하더라도 다른 주제에 사용되는 파라미터에 영향을 주지 않음을 보장합니다.

모델 성장 및 정리

mini-AGI는 만나는 데이터에 따라 용량을 동적으로 성장 및 축소하도록 설계되었습니다:

  • 재조합을 통한 성장: 모델이 용량이 부족할 경우, 기존 전문가의 은닉 유닛을 재조합하여 새로운 전문가를 생성합니다. 이는 새로운 전문가가 무작위 노이즈가 아니라 유용한 훈련된 구성 요소를 갖게 함으로써 초기 성능을 보장합니다.
  • 사용 기반 정리: 라우터에 의해 오랫동안 선택되지 않은 전문가는 "사망한" 것으로 간주되어 디스크에서 삭제되어 공간을 회수합니다.
  • 성장 제약: 특정 조건(가용 디스크/VRAM 공간, 보류 손실의 안정성 - "정직한 브레이크")을 충족해야만 새로운 전문가가 추가됩니다.

성능 및 확장성

보고 시점 기준으로 모델은 약 318.1백만 자를 읽었으며, 169개의 전문가로 구성되어 있습니다. 8개 주제에 대한 보류 손실은 0.8336 nats/자(1.2026 비트/바이트)였습니다.

데이터 확장 추세

확장 분석 결과, 모델은 건강한 거듭제곱 법칙(L ∝ D^-0.239)을 따릅니다. 저자는 0.80 BPB(비트/바이트)에 도달하기 위해 약 19.2억 바이트의 데이터가 필요하며, 이는 단일 노트북 GPU에서 몇 주 내에 훈련 가능하다고 제안합니다.

커뮤니티의 통찰과 비판

기술 사용자들 사이의 논의는 현재 구현의 잠재력과 한계를 모두 강조합니다:

  • 일반화 vs. 기억: 일부 사용자는 모델이 진정으로 일반화하고 있는지, 아니면 단지 고효율 기억을 수행하고 있는지에 대해 의문을 제기했습니다. 한 사용자는 현재 모델의 비트/바이트(BPB) 성능이 enwik9과 같은 특정 데이터셋에 훈련된 매우 작은 밀집 모델보다 높은(더 나쁜) 수준임을 지적했습니다.
  • 아키텍처의 잠재력: 지속 학습의 특성 덕분에 모델이 "프로토-AGI" 역할을 할 수 있다는 점에 관심이 있으며, 중첩 강화 학습이나 자기 유사 재귀 구조를 탐색할 것을 제안했습니다.
  • 학습률에 대한 우려: 비판자들은 트렁크 학습률 감소가 치명적인 망각을 완전히 제거하는 것이 아니라 단지 지연시킬 뿐이며, 매우 긴 시간 동안 트렁크도 결국 망각을 겪을 것이라고 지적했습니다.

"전문가 풀이 망각을 막는 것은 아닙니다. 작업 세트를 고정하는 것... 단지 0.3571 nats만 비용이 듭니다... 대부분의 가중치를 보존하는 것은 충분하지 않습니다."

기술 사양 요약

구성 요소 세부 사항
어휘 265(256바이트 + 9개 마커)
컨텍스트 창 4,096
핵심 아키텍처 RMSNorm, RoPE, SwiGLU, Flash Attention
깊이 2개의 프리루드 블록 + 1개의 순환 블록(최대 24회 적용)
VRAM 요구 사항 최소 8GB
파라미터 관리 페이지 기반 MoE(디스크 $
ightarrow$ RAM $
ightarrow$ VRAM)

Sources

관련