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

2 weeks ago 9

  • 최신 언어 모델(LLM) 은 복잡한 수학 문제를 풀 수 있지만, 단순한 계산 작업에서는 긴 추론 단계와 문맥 길이 때문에 실패율이 높음
  • 연구팀은 트랜스포머 내부에 실제 컴퓨터를 구현, 외부 도구 없이 C 코드나 WebAssembly 프로그램을 모델이 직접 실행하도록 설계
  • 핵심 기술은 2차원(2D) 헤드 기반의 빠른 디코딩 경로, 이를 통해 선형 시간 대신 로그 시간으로 어텐션 조회를 수행
  • 이 방식으로 헝가리안 알고리듬이나 Sudoku 풀이기 같은 복잡한 프로그램을 모델 내부에서 수백만 단계까지 정확히 실행 가능
  • 결과적으로 트랜스포머가 단순한 언어 모델을 넘어 실제 계산을 수행하는 컴퓨터로 기능할 수 있음을 입증

LLM은 왜 계산을 못하는가

  • 최신 LLM은 국제수학올림피아드 수준의 문제를 풀 수 있지만, 덧셈·Sudoku 같은 단순 계산에는 실패율이 높음
    • Sudoku-Bench 등 벤치마크에서 외부 도구 없이 해결률이 낮음
  • 현재는 두 가지 우회 방식이 사용됨
    • 도구 사용(tool use): 모델이 코드를 생성하고 외부 인터프리터가 실행
    • 에이전트 오케스트레이션(agentic orchestration): 외부 루프가 상태를 저장하고 모델을 반복 호출
  • 그러나 이런 방식은 계산이 모델 외부에서 수행된다는 한계를 가짐
  • 진정한 목표는 모델이 내부적으로 계산을 실행해 스스로 컴퓨터처럼 동작하는 것임

트랜스포머를 컴퓨터로 만드는 방법

  • 트랜스포머는 이론적으로 튜링 완전성을 가지지만, 실제로는 효율적인 실행이 어려움
  • 연구팀은 RAM 컴퓨터를 트랜스포머 내부에 구현, 각 명령어를 최대 5개의 토큰으로 매핑
  • 문제는 디코딩 과정의 비효율성으로, 일반 트랜스포머는 매 단계마다 전체 히스토리에 주의를 기울여야 함
  • 이를 해결하기 위해 헤드 차원을 2로 제한하고, 로그 시간 조회가 가능한 디코딩 경로를 설계
    • 이로써 수백만 단계의 프로그램을 단일 트랜스포머 실행으로 처리 가능

모델 내부 실행(In-model Execution)

  • 기존 LLM은 3 + 5 계산 시 코드를 생성하고 외부 인터프리터가 결과를 반환
  • 본 시스템은 WebAssembly 인터프리터를 트랜스포머 가중치 내부에 구현, 모델이 직접 명령어를 실행
    • 예: i32.const 03, i32.add, output 등의 명령을 토큰으로 생성하며 실행 추적(trace)을 출력
  • 실행 과정의 모든 중간 상태가 토큰 스트림으로 투명하게 출력, 외부 호출 없이 모델 내부에서 완료

Sudoku 실행 데모

  • Sudoku는 긴 제약 조건을 가진 계산 문제로, 기존 신경망은 어려움을 겪음
  • 본 시스템은 컴파일된 Sudoku 솔버를 트랜스포머 내부에서 직접 실행, 100% 정확도 달성
    • Arto Inkala의 ‘가장 어려운 Sudoku’도 3분 이내 해결
  • 정확성은 솔버 코드의 정확성에 의해 보장, 학습된 휴리스틱이 아님

계산을 토큰 추적으로 인코딩하기

  • 트랜스포머는 수정 불가능한 과거 토큰을 참조하며 새 토큰을 생성
  • 계산을 “추가만 가능한(trace that only appends)” 형태로 표현
    • 각 단계는 이전 몇 개의 위치만 참조해 다음 상태를 기록
  • 예시: 문장 내 동사의 개수를 세는 알고리듬은 두 위치(입력 단어, 이전 상태)만 참조
  • 모델은 이러한 추적을 통해 가상 머신의 상태(명령 포인터, 스택, 메모리 등) 를 재구성

지수적으로 빠른 어텐션 (Exponentially Fast Attention)

  • 일반 트랜스포머는 시퀀스 길이에 비례해 선형 시간이 소요되어 긴 실행 추적에 비효율적
  • 제안된 HullKVCache는 조회를 O(log t) 시간에 수행, 기존 KVCache 대비 수천 배 빠름
    • 예시: 동일 작업에서 23,513 tok/s vs 62 tok/s
  • 긴 실행 추적에서도 일정한 처리 속도를 유지, 수백만 단계의 계산 실행 가능

2D 어텐션의 기하학적 원리

  • 각 키(key)는 2차원 벡터로 표현되며, 쿼리(query)는 평면상의 방향으로 해석
  • 가장 큰 내적값을 갖는 점을 찾는 문제는 볼록 껍질(convex hull) 상의 점 탐색으로 변환
  • 이 구조를 유지하면 로그 시간 내 조회 가능
  • 인덱스 조회를 위해 키를 (2j, -j²)로 설정하면, 쿼리 (i, 1)일 때 정확히 i번째 인덱스를 선택
  • 누적합 연산도 지원되어 명령 포인터·스택 깊이 추적 가능

향후 발전 방향

  • k-희소 소프트맥스(k-sparse softmax) 를 통해 근사적 소프트맥스 어텐션 구현 가능
    • 복잡도는 O(k + log n)
  • 2D 헤드 트랜스포머는 동일한 파라미터 예산으로도 다수의 헤드를 활용 가능
    • 예: d_model=256, n_heads=128, d_h=2
  • 결합형 시스템 구상
    • 빠른 계산용 2D 트랜스포머와 일반 LLM을 결합
    • LLM이 계획·추론을 담당하고, 2D 트랜스포머가 정확 계산 실행
  • 프로그램을 가중치로 컴파일 가능
    • C 코드 → 트랜스포머 가중치(WQ, WK, WV, WO, Wff)
    • 학습 없이 프로그램 로직을 직접 내장
  • 장기적으로는 AI 시스템이 소프트웨어처럼 성장, 새로운 계산 모듈을 내부에 추가 가능

결론

  • 트랜스포머가 외부 도구 없이 자체 추론 루프 내에서 프로그램을 실행할 수 있음을 실증
  • 이는 학습된 표현과 컴파일된 알고리듬의 통합을 가능하게 하며,
    의료·공급망·금융 등 복잡한 의사결정 문제 해결에 필요한 정확하고 신뢰할 수 있는 계산 능력을 제공
  • 연구팀은 이러한 계산 가능한 AI 시스템을 구축 중임

Read Entire Article