포스트

LLM을 컴퓨터로 만들기 - 트랜스포머 내부에서 프로그램을 실행하는 방법

목차

  1. 개요
  2. 배경: LLM의 계산 한계
  3. 방법론
  4. 주요 결과
  5. 향후 방향
  6. 결론
  7. Reference

개요

언어 모델은 연구 수준의 수학 문제를 풀 수 있지만, 다수의 단계와 긴 컨텍스트를 필요로 하는 단순 계산 작업에서 어려움을 겪는다. 두 수의 곱셈이나 간단한 스도쿠도 외부 도구 없이는 거의 불가능하다. 이 연구는 트랜스포머 내부에 컴퓨터를 구축하여, 임의의 C 프로그램을 토큰으로 변환하고 모델 자체가 수백만 단계를 몇 초 만에 실행할 수 있게 만드는 방법을 제안한다.

배경: LLM의 계산 한계

최첨단 LLM은 국제수학올림피아드 금메달 수준의 수학을 풀 수 있지만, 순수 계산 작업에서는 기본 덧셈도 실수한다. 현재 이를 우회하는 두 가지 방법이 있다.

첫 번째는 도구 사용(Tool use)이다. 모델이 코드를 작성하고 외부 인터프리터가 이를 실행하는 방식이다.

두 번째는 에이전트 오케스트레이션이다. 외부 루프가 중간 상태를 저장하고 모델을 반복 호출하는 방식이다.

이를 비유하면, 인간은 날 수 없다. 비행기를 만들었다고 인간이 나는 것은 아니다. 마찬가지로 LLM에 외부 도구를 붙여도 모델 자체가 계산하는 것은 아니다.

진짜 질문은 다음과 같다. 모델이 내부적으로 계산을 실행할 수 있는가? 안정적으로, 효율적으로, 매우 긴 시간 동안 가능한가?

방법론

트랜스포머 아키텍처는 이론적으로 튜링 머신을 시뮬레이션할 수 있음이 증명되어 있다. 하지만 이론적 보편성과 실제 실행 효율은 다른 문제이다. 이 연구에서는 순수 이론적 계산 모델 대신 현대적 RAM 컴퓨터를 트랜스포머 내부에 구현하는 접근법을 택했다.

WebAssembly 인터프리터 구현

핵심 아이디어는 WebAssembly 인터프리터를 트랜스포머 가중치 안에 구현하는 것이다. WebAssembly는 C/C++에서 컴파일 가능한 저수준 명령어 세트이다. 각 명령어가 최대 5개 토큰에 매핑된다.

도구 사용(Tool use)과 모델 내부 실행(In-model execution)의 차이는 다음과 같다.

도구 사용 방식에서는 모델이 코드를 외부에 보내고 블랙박스 답변을 받는다. 이 과정은 불투명하다.

반면 모델 내부 실행 방식에서는 모든 중간 단계가 트레이스에 나타난다. 모델이 자체 디코딩 루프를 떠나지 않으며, 전 과정이 투명하다.

계산의 트레이스 인코딩

계산을 “추가만 가능한 트레이스”로 인코딩한다. 노트북처럼 모든 계산 단계가 다음 줄에 기록된다. 한번 기록된 줄은 변경 불가하며, 성장만 가능하다. 각 단계에서 이전 줄을 참조할 수 있지만 편집은 불가능하다. 이것이 오토리그레시브 트랜스포머의 작동 방식과 유사하다.

2D 어텐션과 볼록 껍질 기반 빠른 디코딩

표준 트랜스포머 디코딩에서는 각 단계마다 전체 프리픽스에 대해 어텐션을 수행한다. t번째 단계에서 t에 비례하는 작업량이 필요하며, 전체 비용은 이차적으로 증가한다.

이 문제를 해결하기 위해 2D 어텐션 헤드를 도입했다. 헤드 차원을 2로 제한한다(d_model = 36, n_heads = 18이면 헤드당 2D). 2D 키는 평면 위의 점이 되고, 쿼리는 방향 벡터가 된다. 최대 내적을 구하는 것은 기하학에서 “볼록 껍질(convex hull)에서 가장 먼 점 찾기” 문제가 된다. 볼록 껍질 데이터 구조를 사용하면 기존 O(t) 대비 O(log t) 시간에 검색이 가능하다.

모델은 완전한 표준 PyTorch 트랜스포머이다.

1
2
3
class VanillaTransformer(nn.Module):
    def __init__(self, vocab, d_model=36, n_heads=18, n_layers=7, d_ffn=36):
        ...

커스텀 어텐션 커널이나 스파스 마스크는 사용하지 않는다. 가중치만 특별하다.

주요 결과

성능 비교

HullKVCache와 표준 KVCache의 성능 비교는 다음과 같다.

방식처리 속도생성 토큰 수소요 시간
HullKVCache31,037 tok/s41,7091.3초
표준 KVCache379 tok/s13,96736.9초 시점

약 80배 이상의 속도 차이를 보였다.

스도쿠 데모

세계에서 가장 어려운 것으로 알려진 Arto Inkala의 스도쿠를 3분 이내에 해결했다. 완전한 컴파일된 스도쿠 솔버를 트랜스포머 내부에서 실행한 결과이다. 학습된 휴리스틱이 아닌 실제 알고리즘을 실행한 것이다. 이를 통해 오토리그레시브 패러다임 자체가 아니라, 긴 실행 트레이스에서의 표준 어텐션 비용이 진짜 병목임을 증명했다.

향후 방향

이 연구에서 제시하는 향후 방향은 다음과 같다.

첫째, 더 풍부한 어텐션 메커니즘이다. k-sparse softmax로 O(k + log n) 비용 달성이 가능하다.

둘째, 2D 헤드로 대규모 모델 학습이다. 헤드와 레이어 수를 늘려 파라미터 수를 유지할 수 있다.

셋째, 하이브리드 시스템이다. LLM이 계획과 추론을 담당하고, 실행 컴포넌트가 알고리즘을 실행하는 구조이다.

넷째, 프로그램을 가중치로 컴파일하는 것이다. 임의의 프로그램을 트랜스포머 가중치로 직접 컴파일할 수 있다.

다섯째, 소프트웨어처럼 성장하는 AI 시스템이다. 새로운 계산 능력을 모델 내부 실행 엔진에 점진적으로 추가하는 방식이다.

결론

이 연구는 트랜스포머가 자체 추론 루프 내에서 프로그램을 효율적으로 실행할 수 있음을 증명했다. WebAssembly 인터프리터를 트랜스포머 가중치에 구현하고, 2D 어텐션 헤드와 볼록 껍질 데이터 구조를 활용하여 지수적으로 빠른 디코딩을 달성했다. 미래 AI 시스템은 소프트웨어를 사용하는 것이 아니라 소프트웨어를 포함하게 될 것이다.

Reference