트랜스포머 내부에서 프로그램을 실행해 지수적으로 빠른 추론 달성

- LLM이 수학 올림피아드 수준의 문제를 풀면서도 단순한 덧셈/스도쿠조차 외부 도구 없이는 정확히 수행하지 못하는 한계를 극복하기 위해, 트랜스포머 내부에 실제 컴퓨터를 구축하는 접근법 - 임의의 C 코드를 토큰으로 변환하여 모델 자체가 수백만 스텝의 실행 트레이스를 직접 생성하며, CPU에서 초당 3만 토큰 이상의 속도로 스트리밍 가능 - 핵심 기술은 어텐션 헤드 차원을 2D로 제한하여 볼록 껍질(convex hull) 기반의 기하학적 검색으로 전환, 선형 시간 스캔을 로그 시간으로 대체 - WebAssembly 인터프리터를 트랜스포머 가중치에 구현하여, 외부 도구 호출 없이 모델의 디코딩 루프 안에서 프로그램 실행이 투명하게 이루어짐 - 향후 프로그램을 가중치로 직접 컴파일하고, 학습된 표현과 컴파일된 알고리듬을 단일 연산 기반에 통합하는 AI 시스템으로의 확장 가능성 제시 LLM의 연산 한계 - 최첨단 LLM은 국제수학올림피아드 금메달 수준의 문제를 풀거나 미해결 수학·과학 문제에 도전할 수 있지만, 순수 연산 작업에서는 여전히 실패 - 기본적인 덧셈에서도 오류가 발생하고, Sudoku-Bench 같은 벤치마크에서 외부 도움 없는 풀이 성공률이 매우 낮음 - 현재 이 격차를 메우는 두 가지 우회 방법 존재 - 도구 사용(Tool use): 모델이 코드를 작성하면 외부 인터프리터가 실행하고 결과를 반환 - 에이전트 오케스트레이션: 외부 루프가 중간 상태를 저장하고 작업을 분해하여 모델을 반복 호출 - 이러한 접근은 유용하지만, 실제 연산 능력이 모델 외부에 존재한다는 근본적 한계를 부각 - 비유하면 인간이 비행기를 만들었다고 해서 인간이 날 수 있는 것이 아닌 것과 동일한 상황 - 모델이 연산에 대해 추론하거나 도구를 조율할 수 있지만, 스스로 실행할 수 없으면 진정한 컴퓨터가 아님 트랜스포머를 컴퓨터로 만든 방법 - 트랜스포머 아키텍처가 튜링 머신을 시뮬레이션할 수 있다는 이론적 보편성은 여러 연구에서 입증되어 있으나, 이론적 표현력이 실용적 실행 효율을 보장하지는 않음 - 순수 이론적 계산 모델 대신 현대적 RAM 컴퓨터를 트랜스포머 내부에 구현, 각 명령어가 최대 5개 토큰에 매핑 - 그러나 더 깊은 문제는 디코딩 과정 자체에 존재 - 표준 오토레그레시브 디코딩은 각 스텝에서 계속 늘어나는 전체 히스토리와 상호작용 - KV 캐싱이 과거 키/값 재계산을 피하지만, 캐시 크기에 비례하는 어텐션 비용은 여전히 발생 - 이 한계를 해결하기 위해 어텐션 헤드 차원을 2D로 제한, 실행 스타일 트레이스에 대한 효율적 디코딩 경로를 도입 - 주요 검색/업데이트 연산이 시퀀스 길이에 대해 로그 시간으로 수행 가능 - 이를 통해 수백만 스텝의 프로그램을 트랜스포머 내부에서 실행 가능 연산의 의미 — 도구 사용 vs 모델 내 실행 - 기존 도구 사용 방식: 모델이 python -c "print(3+5)" 같은 코드를 출력 → 외부 인터프리터가 실행 → 결과를 토큰 스트림에 주입 - 이 시스템의 방식: WebAssembly 인터프리터를 트랜스포머 가중치에 구현 - WebAssembly는 빠르고 결정론적 실행을 위한 저수준 명령어 세트이자 C/C++의 범용 컴파일 타겟 - 3 + 5 계산 시 모델이 wasm 명령어( i32.const , i32.add )를 출력한 후 빠른 디코딩 모드로 전환하여 스텝별 실행 트레이스를 직접 생성 - 핵심 차이점 - 도구 사용은 불투명(opaque): 모델이 제어를 넘기고 블랙박스 답변을 수신 - 모델 내 실행은 투명(transparent): 모든 중간 단계가 트레이스에 나타나며, 모델이 자체 디코딩 루프를 벗어나지 않음 스도쿠 데모 — 정확한 장기 연산의 스트레스 테스트 - 학습된 신경망 접근법은 쉬운 스도쿠에서는 좋은 성능을 보이나 어려운 문제에서 완전히 실패 - 오토레그레시브 모델이 제약 만족 문제에 근본적으로 부적합하다는 일반적 설명이 있으나, 이 시스템은 완전히 오토레그레시브이면서도 100% 정확도 달성 - 실제 병목은 오토레그레시브 패러다임 자체가 아니라, 어려운 스도쿠가 매우 긴 실행 트레이스를 요구하고 표준 어텐션이 긴 컨텍스트 생성을 비실용적으로 만드는 것 - 컴파일된 완전한 스도쿠 솔버를 트랜스포머 내부에서 실행, 학습된 휴리스틱이 아닌 실제 알고리듬을 스텝별로 수행 - Arto Inkala의 "세계에서 가장 어려운 스도쿠"를 3분 이내에 정확히 해결 - 컴파일된 솔버가 정확하면 트랜스포머의 실행도 정확하다는 보편적 보장 제공 연산의 인코딩 방식 — 추가만 가능한 트레이스 - 오토레그레시브 트랜스포머를 자신의 히스토리 안에 사는 기계로 비유 - 전통적 컴퓨터는 편집 가능한 메모리를 업데이트하지만, 트랜스포머에는 그런 개념이 없음 - 고정된 프롬프트(입력/프로그램)와 계속 커지기만 하는 트레이스(생성된 토큰)로 구성 - 추가만 가능한 노트북 비유 - 첫 줄들은 입력(프롬프트), 이후 각 줄은 연산의 다음 스텝을 기록 - 이전 줄을 수정할 수 없고, 각 스텝에서 소수의 이전 위치만 참조 가능 - 많은 알고리듬이 "각 스텝에서 고정된 소수의 이전 위치만 참조하는, 추가만 가능한 트레이스"로 표현 가능 - 이 시스템에서 모델이 생성하는 토큰은 가상 머신의 명령어 포인터, 메모리/스택 연산, 산술, 제어 흐름, 출력 등 진화하는 상태를 표현 - 어텐션 헤드는 공유된 1차원 배열처럼 작동하며, 각 토큰이 하나의 인덱스에 값을 쓰고 다른 인덱스에서 값을 읽는 쓰기 후 읽기 프리미티브 제공 - 어텐션이 누적 합산(cumulative sums) 도 계산 가능하여 명령어 포인터, 스택 깊이 등을 델타 증분의 누적 합으로 추적 핵심 잠금 해제: 지수적으로 빠른 어텐션 - 실제 컴퓨터는 명령어당 거의 일정한 작업량으로 소형 상태를 업데이트하지만, 표준 트랜스포머는 t번째 디코딩 스텝에서 길이 t의 프리픽스와 상호작용, 총 비용이 이차적으로 증가 - HullKVCache 도입으로 이 이차적 폭증을 해결 - 실제 벤치마크에서 HullKVCache는 초당 31,037 토큰(6,747줄/초), 표준 KVCache는 초당 272 토큰(59줄/초)으로, 약 114배 차이 - 스텝당 Θ(t) 시간 대신 O(log t) 시간으로 어텐션 검색 수행 - 2D 어텐션의 빠른 경로 - 트랜스포머를 일반적으로 가속하거나 새로운 아키텍처를 도입하는 것이 아니라, 바닐라 트랜스포머에서 헤드 차원을 2D로 제한하는 다루기 쉬운 파라미터화에 집중 - d_model = 36 , n_heads = 18 로 헤드당 정확히 2차원, 7개 레이어, 표준 nn.MultiheadAttention , 게이트드 피드포워드 네트워크 사용 - 커스텀 어텐션 커널이나 희소 마스크 없이 순수 바닐라 PyTorch로 구성 - 레이어 수, 헤드 수, 임베딩 크기는 임의로 크게 설정 가능하므로 전체 모델이 작을 필요 없음 - n_heads = d_model / 2로 더 많은 헤드를 사용 - 2D 어텐션의 기하학적 관점 - 어텐션 메커니즘: 과거 토큰이 2D 키 벡터 kⱼ와 값 vⱼ를 제공, 현재 스텝이 쿼리 q를 형성하여 내적이 최대인 키의 값을 반환 (하드맥스 어텐션) - 이는 계산기하학의 고전적 "지지점(supporting point) 쿼리" 와 정확히 동일: 방향 q가 주어졌을 때 볼록 껍질에서 해당 방향으로 가장 먼 점을 찾는 문제 - 토큰 생성과 동시에 볼록 껍질 자료구조를 유지하여 전체 과거 포인트에 대해 로그 시간 검색 가능 - 메모리/스택 연산을 위해 "인덱스 i에 가장 최근 저장된 값" 조회가 필요하며, 이것이 2D 키를 요구하는 이유 - 인덱스 j를 2D 키 kⱼ = (2j, -j²)로 저장하고 방향 q = (i, 1)로 쿼리하면 이차 항 -j²가 정확한 매칭만 이기도록 페널티 역할 - 튜링 완전성을 위해 2D 어텐션이면 충분하며, 이 블로그에서 보여준 것처럼 전체 RAM 컴퓨터를 표현할 수 있음 향후 계획 - 더 풍부한 어텐션 메커니즘 - 현재 구현은 하드맥스 어텐션 사용이지만 이는 근본적 제약이 아님 - k-sparse 소프트맥스 어텐션으로 근사 가능: 중첩된 볼록 껍질에서 상위 k개 키를 검색하여 그것들에 대해서만 소프트맥스 수행, 비용 O(k + log n) - 기하학적 빠른 경로가 실행기 구조에만 국한되지 않으며, 원칙적으로 2D 헤드를 가진 모든 트랜스포머의 디코딩 시간을 가속 가능 - 3D 헤드로도 3D 볼록 껍질을 통해 자연스럽게 확장 가능하나, 더 높은 차원은 효율이 급격히 감소 - 2D 헤드로 대규모 모델 학습 - 2D 헤드 파라미터화는 본질적으로 작은 것이 아님 — 더 많은 헤드와 레이어를 사용하여 표준 트랜스포머와 비슷한 총 파라미터 수 유지 가능 - 여러 활용 모드 가능 - 더 느리고 범용적인 모델과 결합하는 전용 빠른 경로 - 단일 시스템 내의 빠름/느림 하이브리드 아키텍처 - 2D 모델이 토큰을 빠르게 제안하고 일반 어텐션 모델이 검증하는 추측적 디코딩(speculative decoding) - 실행 트레이스가 포워드 패스의 일부이므로 전체 과정이 미분 가능(differentiable) — 외부 도구와는 근본적으로 다르며, 더 큰 모델에 직접 통합 가능한 학습 가능한 연산 기반 - 프로그램을 가중치로 컴파일 - 현재는 모델이 가중치에 인코딩된 인터프리터를 학습하지만, 가중치 생성 컴파일 기계를 확장하면 임의의 프로그램을 토큰 시퀀스 없이 직접 가중치로 컴파일 가능 - 가중치 자체가 소프트웨어의 배포 대상(deployment target) 이 되어, 모델이 소프트웨어 같은 행동을 학습하는 것이 아니라 컴파일된 프로그램 로직을 내부 회로의 일부로 포함 - 로직을 가중치로 컴파일할 수 있다면, 경사 하강법이 모델을 수정하는 유일한 방법이 아니게 됨 — 구조, 알고리듬, 보장을 네트워크에 직접 삽입하는 또 다른 경로 - 궁극적으로 경험에서 학습할 뿐만 아니라 자신의 가중치를 수정하거나 확장하여 내부 기계를 스스로 재작성하는 시스템으로 발전 가능 - AI 시스템을 소프트웨어처럼 성장시키기 - 현대 소프트웨어 생태계가 모듈, 추상화, 재사용 가능한 컴포넌트를 축적하며 진화하듯, AI 시스템 내부에서도 새로운 연산 능력이 점진적으로 추가되는 유사한 과정이 가능 - 새로운 기능이 전체 시