8GB 노트북 GPU로 멈추지 않고 학습하는 바이트 단위 언어모델 'mini-AGI' 공개

GitHub에 volotat/mini-AGI라는 실험적 프로젝트가 올라왔다. 8GB VRAM을 가진 노트북이나 데스크톱 GPU 한 장에서 처음부터 학습을 시작하고, 이후 읽는 모든 텍스트로 학습을 이어가는 바이트 단위 언어모델이다. 눈에 띄는 건 메모리 운용 방식이다. 가중치를 GPU 메모리에 상주시키는 대신 디스크에 평범한 파일로 저장해 두고, 필요할 때만 카드로 끌어올린다. 그 결과 모델의 파라미터 수는 VRAM 용량이 아니라 남은 디스크 공간에 의해 결정된다. 학습 도중 용량이 부족해지면 새 전문가를 늘리고, 아무도 요청하지 않는 것은 잘라낸다.

구조는 일반적인 트랜스포머 스택과 다르다. 입력은 256개 바이트 값 그대로 들어가므로 토크나이저도, 새 어휘를 맞춰 넣어야 할 데이터 타입도 없다. 두 개의 dense prelude 블록을 지난 뒤 하나의 recurrent 블록을 최대 24번 반복 적용하는데, 매 반복마다 공유 전문가 풀에서 top-8 전문가를 스스로 고른다. 문자 하나가 최대 26번의 블록 적용을 거치는 셈이다. 반복 사이의 잠재 상태는 별도로 디코딩되지 않고 어댑터를 통해 임베딩된 입력과 다시 합쳐진다. 루프가 읽고 있는 텍스트에서 멀어지지 않게 하려는 장치다.

깊이는 문자마다 달라진다. halting head가 각 문자를 모든 행에서 평가하고, 다음 행이 답을 바꾸지 않을 것 같으면 거기서 멈춘다. PonderNet 방식으로 학습 중에는 모든 깊이를 계산해 halting 확률로 가중한다. 쉬운 문자는 한 행으로 끝나고 어려운 문자는 여러 행을 쓴다. 라우팅은 문자 단위가 아니라 블록 적용 단위로 일어나기 때문에 같은 전문가가 서로 다른 깊이에서 여러 번 선택될 수도 있다. 전문가에게 주제가 배정되지 않고 레이블도 없다. soft top-k 라우팅이 스스로 능력을 분배하며, 그 대가로 능력들이 파라미터를 공유해 서로 간섭할 수 있다.

메모리는 세 층으로 나뉜다. 디스크에는 모델이 가진 모든 전문가가 있고, 그 위에 최근 요청된 전문가를 담는 캐시(LRU로 축출)와 실제로 문자가 라우팅될 수 있는 워킹셋이 있다. 다음 청크를 읽기 전에 모델은 앞으로 읽을 텍스트가 무엇을 원하는지 예측하고, 그 답이 워킹셋이 된다. 이 수요는 직전 청크를 읽을 때 호출 지점이 실제로 라우팅에 사용한 은닉 상태를 기준으로 채점된다. 임베딩만 보고 채점하면 맥락이 없어 모든 주제가 같은 전문가를 요구하게 되기 때문이다. Adam 옵티마이저의 모멘트는 VRAM 슬롯이 아니라 전문가에 딸려 다닌다.

배경에는 소유 문제가 있다. 지금 실제로 손에 넣을 수 있는 언어모델은 누군가가 학습해 동결해 놓은 것이다. 주변부를 파인튜닝할 수는 있어도 자기 하드웨어에서 처음부터 학습할 수는 없고, 일상에서 얻은 데이터로 계속 학습시키면 이전에 알던 것을 잊어버린다. 결국 개인 모델은 남의 모델 위에 얇게 얹은 자신일 뿐이고, 배포되는 순간 학습을 멈춘다. 이 프로젝트는 그 지점을 정면으로 겨냥한다.

읽기와 쓰기가 같은 코드 경로를 쓴다. 학습과 생성이 동일한 forward pass이고, 차이는 다음 문자가 파일에서 오는지 모델 자신의 argmax에서 오는지뿐이다. 별도의 파인튜닝 체제도, 동결된 베이스도 없다. 위치 인코딩은 rotary이며 학습 파라미터를 갖지 않기 때문에 컨텍스트 윈도우를 늘리려면 무언가를 재초기화할 필요 없이 학습을 계속하면 된다.

실제 실행에서 관찰된 수치도 공개됐다. 같은 주제를 다룰 때 읽기는 문자당 약 8.0행, 쓰기는 약 9.9행을 쓴다. 워킹셋은 64문자마다 다시 선택된다. 243M 문자를 소화한 시점의 생성 결과는 문법이 맞고 주제에서 벗어나지 않지만 같은 문장을 반복하는 수준이다.

개발자 입장에서 이 실험이 던지는 질문은 명확하다. 학습에 필요한 gradient와 옵티마이저 상태는 가중치의 약 3배에 달하기 때문에, 8GB 카드에 모델을 통째로 올리는 방식은 애초에 성립하지 않는다. 가중치를 디스크로 내리고 워킹셋만 상주시키면 파라미터 상한이 저장 공간으로 바뀐다. 학습이 멈추지 않는 구조라는 점도 실무적으로 의미가 있다. 자기 데이터로, 자기 하드웨어에서, 아무도 끌 수 없는 모델을 돌린다는 구상이 소비자용 카드 한 장에서 출발할 수 있다는 주장이다.

한계는 분명히 밝혀져 있다. 현재는 장난감 수준의 소형 모델이며 최전선 모델의 성능을 기대해서는 안 된다. 단일 데이터 스트림에서 catastrophic forgetting 없이 지속 학습이 가능하고, 그것이 값싼 하드웨어에서도 된다는 것을 보여주는 실험이라는 위치다. 가중치는 아직 공개되지 않았고 코퍼스 1회독이 진행 중이라 공개까지는 몇 주가 남았다. 최종적인 능력은 실제 하드웨어와 데이터의 규모·품질, 그리고 학습에 투입할 수 있는 시간에 의해 결정된다.