기술 요약: 프라임 푸리에 임베딩 (Prime Fourier Embeddings)

문제 정의

표준적인 신경망 임베딩은 정수의 내재적 대수 구조, 특히 모듈로 산술(modular arithmetic)과 관련된 구조를 드러내는 데 종종 실패한다. 중국인의 나머지 정리(Chinese Remainder Theorem, CRT)에 따르면 합성수 $N$에 대한 모듈로 덧셈은 각 소인수에 대한 독립적인 덧셈으로 분해되지만, 표준적으로 학습된 임베딩은 이러한 분해를 자연스럽게 반영하지 못한다. 결과적으로 모듈로 산술 작업에 대해 훈련된 모델들은 "그로킹(grokking)" 현상을 겪는 경우가 많다. 이는 모델이 일반화에 도달하기 전까지 장기간의 암기 단계를 거치는 현상을 의미한다. 이러한 지연은 표현론적 실패에서 기인한다. 즉, 입력값에 명시적인 소수 국소 구조(prime-local structure)가 결여되어 있을 때, 네트워크는 경사 하강법을 통해 이러한 대수적 관계를 처음부터 재구성해야 하며, 이 과정은 데이터 집약적이고 취약하다.

기존의 접근 방식들은 한계가 있다. 연속적 인코딩(예: xVal)은 크기는 보존하지만 주기성을 상실하며, 푸리에 수 임베딩(Fourier Number Embeddings, FoNE)과 같은 주기적 접근 방식은 합성 기반(예: 10진법)을 사용하여 서로 다른 소수 구조(예: 2-adic 및 5-adic)를 공유 채널 내에 뒤섞어 놓음으로써 소수별 정보를 격리하는 데 실패한다.

방법론: 프라임 푸리에 임베딩 (PFE)

저자들은 유리수 $\mathbb{Q}$의 조화 해석(harmonic analysis)에서 유도된 원칙적인 임베딩 체계인 **프라임 푸리에 임베딩(PFE)**을 제안한다.

• 구성: PFE는 정수 $a$를 소수 인덱스 기반의 $(\cos, \sin)$ 쌍들의 집합으로 인코딩한다. 소수 $p$와 깊이 $d$ ($0 \le d < D$)에 대하여, 임베딩은 다음과 같이 정의된다:
  $$\text{PFE}_{p,d}(a) = \left[ \cos\left(\frac{2\pi a}{p^{d+1}}\right), \sin\left(\frac{2\pi a}{p^{d+1}}\right) \right] \in \mathbb{R}^2$$
  이는 군 $\mathbb{Z}/p^{d+1}\mathbb{Z}$의 캐릭터 $\chi_1$의 실수 분해에 해당한다.
• 이론적 기초: 이 설계는 폰트랴긴 쌍대성(Pontryagin duality)과 아델릭 캐릭터(adelic characters)의 분해에 의해 동기 부여되었다. 이 임베딩은 로컬 $p$-adic 캐릭터를 직접 구현함으로써, 어떤 정수의 잔여 구조(residue structure)라도 소수 간의 결합 없이 해당 소수 채널로부터 읽힐 수 있도록 보장한다.
• 구현: 실험에서 저자들은 16개의 고정된 소수 집합($P_0 = \{3, \dots, 59\}$)과 깊이 제한 $D=3$을 사용하였으며, 그 결과 임베딩 차원은 96이 된다. 임베딩 특징은 결정론적이며 학습되지 않는다(non-trainable).

핵심 이론적 기여

본 논문은 PFE 상에서 작동하는 임의의 선형 사상에 대한 엄격한 구조적 보장을 확립한다:

1. 블록-대각 분해 정리 (정리 3.1): PFE에 대한 곱군 작용(product group action)에 대해 등변(equivariant)인 임의의 선형 사상 $W$는 반드시 블록-대각(block-diagonal) 형태여야 한다.
   • 각 소수 $p$와 깊이 $d$마다 하나의 독립적인 블록이 존재한다.
   • 서로 다른 소수 간의 블록 사이에는 결합이 전혀 없다(zero coupling).
   • 동일한 소수 내에서 서로 다른 깊이 간의 블록 사이에도 결합이 전혀 없다.
2. 증명 메커니즘: 이 결과는 슈어의 보조정리(Schur's Lemma)를 따른다. 각 PFE 블록은 가법군 $\mathbb{Z}$의 비동형 기약 표현(non-isomorphic irreducible representation, 즉 캐릭터)을 나타낸다. 서로 다른 $(p, d)$ 쌍에 대한 캐릭터들은 서로 동형이 아니므로, 이들 사이의 임의의 등변 인터티너(intertwiner)는 0이어야 한다.
3. CRT 예측: 모듈러스 $N = p_1 \cdots p_k$인 작업의 경우, 이 이론은 $N$의 인수인 $k$개의 소수 채널만이 작업 관련 정보를 운반할 것이며, 나머지 모든 채널은 무관할 것이라고 예측한다.

실험 결과

저자들은 모듈로 덧셈 작업 $(a+b) \mod N$에 대한 체계적인 절제 연구(ablation studies)를 통해 이러한 예측을 검증한다.

• 실험 1 (단일 소수 작업): 다양한 소수 모듈러스를 가진 작업에 대해 모델을 훈련했다. 특정 소수 행(task modulus에 해당하는 행)을 절제(zeroing)했을 때 테스트 정확도가 크게 떨어졌다. 반면, 무관한 소수 행을 절제했을 때는 정확도 저하가 거의 없었다.
  • 결과: 대부분의 구성에서 특수화 비율(대각 성분 감소 대 비대각 성분 감소)이 500배 이상으로 나타났으며, 이는 네트워크가 오직 관련 소수 채널만을 활용함을 확인시켜 준다.
• 실험 2 (합성 모듈러스): 제곱 자유(square-free) 합성 모듈러스(예: $N=15, 21, 105$)에 대해 모델을 훈련했다.
  • 결과: $N$의 소인수를 절제했을 때는 큰 정확도 저하가 발생했으나, $N$의 인수가 아닌 것을 절제했을 때는 미미한 저하만을 보였다. 인수인 경우와 아닌 경우의 비율 역시 500배 이상이었다.
  • 정확도: 모델은 테스트된 모든 제곱 자유 합성 모듈러스 및 입력 범위(입력 범위가 $N$에 비해 충분히 큰 경우)에 대해 100%의 분포 내(in-distribution) 테스트 정확도를 달랐다.

의의 및 주장

본 논문은 PFE가 모듈로 산술의 학습을 **표현 발견 문제(representational discovery problem)**에서 **선택 문제(selection problem)**로 전환한다고 주장한다.

• 원칙적 vs 임의적: 휴리스틱한 주파수 선택과 달리, 소수 기저(prime basis)는 $\mathbb{Q}$의 아델릭 푸리에 분해에 수학적으로 근거한다. 이는 초기화 단계부터 필요한 대수적 구조가 존재하도록 보장한다.
• 기계적 해석 가능성: 본 연구는 임베딩 설계가 표현하고자 하는 대상의 수학에 의해 가이드될 수 있음을 보여준다. 네트워크의 내부 구조는 훈련 전(표현론을 통해)에 특징지어질 수 있으며, 훈련 후(표적 절제를 통해) 검증될 수 있다.
• 반증 가능한 이론: 블록-대각 분해 정리는 네트워크 행동에 대한 증명 가능한 구조적 제약을 제공한다. 이 제약에 대한 경험적 확인은 사후 관찰을 넘어 수학적으로 근거 있는 이론을 통해, 신경망의 더 신뢰할 수 있고 투명한 수치 추론을 향한 경로를 제시한다.

저자들은 블록-대각 구조는 증명되었으나, 경사 하강법이 구체적으로 어떤 메커니즘을 통해 대칭을 깨뜨리지 않고 등변 해(equivariant solution)로 수렴하는지는 향후 연구 과제로 남는 열린 문제라고 언급하였다.