Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer

1. 문제 정의 (Problem)

• 배경: 최근 대규모 언어 모델 (LLM) 은 코드 생성 분야에서 뛰어난 성과를 보였으나, 이는 주로 인간이 작성한 데이터를 학습한 결과입니다. 반면, '알고리즘 추출 (Algorithm Extraction)' 은 인간이 작성한 코드에 의존하지 않고, 알고리즘 태스크로 훈련된 모델로부터 실행 가능한 프로그램을 직접 합성하여 새로운 알고리즘을 발견하려는 패러다임입니다.
• 현재의 한계: 기존 연구 (예: MIPS) 는 RNN(순환 신경망) 기반 모델에서 기계적 해석성 (Mechanistic Interpretability) 을 통해 알고리즘 추출에 성공했습니다. 그러나 현재 주류인 Transformer 아키텍처로 이 패러다임을 확장하는 데는 근본적인 장벽이 존재합니다.
  • 표현의 얽힘 (Representation Entanglement): Transformer 는 '중첩 (Superposition)' 현상을 보이며, 특징들이 개별 뉴런이 아닌 겹치는 방향 (비직교 방향) 으로 인코딩됩니다. 이로 인해 내부 표현이 고도로 얽히고 다의적 (Polysemantic) 이 되어, 심볼릭 회귀 (Symbolic Regression) 와 같은 기법을 통해 명시적인 심볼 식을 추출하는 것이 거의 불가능합니다.
• 핵심 질문: Transformer 의 가중치로부터 실행 가능하고 해석 가능한 프로그램을 추출할 수 있는가?

2. 제안 방법론: Discrete Transformer (Methodology)

저자들은 연속적인 표현과 이산적인 심볼 논리 사이의 간극을 메우기 위해 Discrete Transformer라는 새로운 아키텍처를 제안했습니다.

A. 아키텍처 설계 (RASP 모델 정렬)

기존 Transformer Programs 와 RASP(Restricted Access Sequence Processing) 계산 모델을 기반으로 하여, 기능적 분리 (Functional Disentanglement) 를 강제합니다.

1. Numerical Residual Stream (수치적 잔류 스트림):
   • 기존 Transformer 의 밀집 벡터 대신, 명시적인 스칼라 변수들의 집합으로 잔류 스트림을 정의합니다.
   • 이를 통해 정보 간섭을 줄이고, 계산 이력을 분리된 좌표로 유지합니다.
2. Numerical Attention (라우터 역할):
   • 특징을 혼합하는 것이 아니라, 정보 라우팅 (Information Routing) 전용으로 설계됩니다.
   • Piecewise Linear Encoding (PLE): 스칼라 입력을 고차원 공간으로 매핑하여 어텐션 점수를 계산합니다.
   • Hard Attention: 샘플링 함수를 통해 어텐션 패턴을 결정론적인 포인터 (예: "이전 토큰 복사") 로 수렴시킵니다.
3. Numerical MLP (연산자 역할):
   • 산술 및 논리 변환을 담당합니다.
   • 서브-MLP 분해: 복잡한 MLP 를 독립적인 작은 서브 모듈들로 분해하며, 각 모듈은 소수의 입력 (보통 $k=2$) 을 받아 단순한 산술 연산을 수행하도록 설계됩니다. 이는 심볼릭 회귀가 수행하기 쉬운 구조를 만듭니다.
4. Linear Output Head:
   • 추출된 원시 (Primitives) 를 최종 예측으로 통합합니다.

B. 학습 전략: Temperature Annealing

• 연속적 최적화에서 이산적 논리로: 각 모듈에 온도 제어 샘플링 (Temperature-controlled Sampling) 메커니즘을 도입합니다.
• 학습 초기에는 높은 온도 ($\tau$) 를 유지하여 탐색 (Exploration) 을 허용하고, 학습이 진행됨에 따라 온도를 서서히 낮추어 (Annealing) 모델이 완전히 이산적인 계산 그래프로 수렴하도록 유도합니다.

C. 알고리즘 추출 파이프라인 (From Weights to Code)

학습이 완료된 Discrete Transformer 는 다음과 같은 모듈식 전략으로 코드로 변환됩니다.

1. Numerical Attention (가설 검정):
   • 어텐션 패턴을 분석하여 고정 오프셋 (Fixed Offset) 또는 윈도우 극값 (Windowed Extrema) 과 같은 해석 가능한 라우팅 패턴을 통계적 가설 검정으로 식별합니다.
2. Numerical MLP (심볼릭 회귀):
   • 각 서브-MLP 의 입력 - 출력 쌍을 수집하여 PySR과 같은 심볼릭 회귀 도구를 사용하여 닫힌 형식의 수학적 식 (Closed-form expression) 을 도출합니다.
3. 전체 프로그램 조립:
   • 추출된 라우팅 규칙과 산술 식을 결합하여 최종 Python 코드를 생성합니다.

3. 주요 기여 (Key Contributions)

1. Discrete Transformer 아키텍처 제안: RASP 와 정렬된 아키텍처를 통해 연속 최적화와 이산 논리 간의 기능적 분리를 강제하여, Transformer 에서 알고리즘 추출을 가능하게 했습니다.
2. 모듈식 추출 파이프라인 개발: 가설 검정과 심볼릭 회귀를 결합하여, 연속 변수를 다루는 기계적 해석성 기반 프로그램 합성 범위를 확장했습니다.
3. 강건한 분석 및 제어 가능성: 아키텍처적 제약 (Inductive Biases) 이 알고리즘 발견을 어떻게 제어하는지 엄밀하게 분석했습니다. 특히, 기능적 수렴이 구조적 이산화보다 먼저 발생한다는 동역학적 특성을 규명했습니다.

4. 실험 결과 (Results)

다양한 알고리즘 태스크 (선형 산술, 비선형 조합, 연속 동역학) 에서 Discrete Transformer 를 평가했습니다.

• 성능: RNN 기반의 기존 방법 (MIPS) 과 comparable 한 예측 정확도 (MSE $\approx$ 0, Accuracy 1.00) 를 달성했습니다.
• 추출 정확도: 추출된 프로그램은 대부분의 태스크에서 1.00 의 정확도를 보이며, 인간이 읽을 수 있는 Python 코드로 변환되었습니다.
  • 예시: parity_last2 태스크에서 추출된 식 $y_t = x_t + x_{t-1} - 2x_t x_{t-1}$ 은 논리 XOR 연산의 대수적 동치식을 정확히 복원했습니다.
• 연속 동역학 처리: 기존 MIPS 는 이산적 잠재 변수에 제한되어 연속 값 (부동 소수점) 을 처리하지 못했으나, Discrete Transformer 는 연속 변수를 네이티브로 처리하여 지수 이동 평균 (Exponential Moving Average) 및 자유 낙하 높이 (Free Fall Height) 와 같은 태스크에서도 성공적으로 알고리즘을 추출했습니다.
• 제어 가능성 (Controllability): 아키텍처 제약 (예: MLP 개수 조절, 정보 흐름 마스킹) 을 변경함으로써 모델이 동일한 태스크에 대해 서로 다른 논리적 경로 (예: MLP 계산 대신 Attention 을 통한 직접 복사) 를 발견하도록 유도할 수 있음을 입증했습니다.

5. 의의 및 결론 (Significance)

• Transformer 해석성 혁신: Transformer 의 '블랙박스' 내부에서 숨겨진 알고리즘 논리를 명시적인 코드로 복원할 수 있는 첫 번째 체계적인 프레임워크를 제시했습니다.
• 새로운 알고리즘 발견: 인간이 작성한 데이터에 의존하지 않고, 모델이 데이터 분포에서 새로운 알고리즘을 스스로 발견 (De novo discovery) 할 수 있음을 보여줍니다.
• 과학적 발견 도구: 물리 법칙이나 복잡한 수학적 관계를 가진 데이터에서 해석 가능한 공식을 자동으로 추출할 수 있는 강력한 도구로서, 기계 학습과 과학적 발견 (Scientific Discovery) 의 교차점을 개척했습니다.

요약하자면, 이 논문은 Transformer 의 표현력 한계를 극복하기 위해 이산적 구조를 아키텍처 수준에서 강제하고, 이를 통해 학습된 가중치를 실행 가능한 인간 친화적 코드로 변환하는 새로운 패러다임을 제시했습니다.