Prime Fourier Embeddings: A Principled Basis for Modular Arithmetic

Motivation

• Known: 모듈러 연산(modular arithmetic)은 CRT에 의해 소인수별로 독립적인 덧셈으로 분해될 수 있으며, xVal, FoNE, 삼각함수 기반 tabular embedding 등 주기적 구조를 노출하는 다양한 신경망 임베딩 기법들이 제안되어 왔다.
• Gap: 기존의 FoNE 등 주기적 임베딩은 base-10 같은 임의의 주파수 기저를 사용하여 서로 다른 2-adic, 5-adic 구조가 하나의 채널에 뒤섞이는 문제가 있고, 표준 학습 임베딩은 모듈러 구조를 gradient descent로 처음부터 복원해야 하므로 데이터 비효율적이고 취약하다.
• Why: 모델이 학습을 통해 발견해야 할 대수적 구조를 애초에 입력 표현 자체에 내재시킴으로써, 모듈러 연산 학습 문제를 구조 발견이 아닌 채널 선택 문제로 단순화할 수 있다는 원리적 설계 지침을 제공한다.
• Approach: Q의 조화해석(harmonic analysis)과 adelic 문자(character) 분해에 기반해 각 소수 p마다 독립적인 (cos, sin) 쌍을 부여하는 PFE를 구성하고, product group action에 대해 equivariant한 선형 사상은 반드시 소수별 block-diagonal이어야 함을 Schur's lemma로 증명한 뒤, ablation 실험으로 CRT 기반 채널 관련성 예측을 검증한다.

Achievement

Figure 4. Experiment 1. Mean diagonal drop (ablate own prime).

1. PFE 구성 제안: 정수를 소수 p, 깊이 d마다 [cos(2πa/p^{d+1}), sin(2πa/p^{d+1})]로 인코딩하여 Z/p^{d+1}Z의 문자(character) χ1의 실수부·허수부로 해석되는 사전 구조화된(pre-structured) 표현을 제시.
2. Block-Diagonal Decomposition Theorem 증명: PFE의 product group action에 equivariant한 임의의 선형 사상은 소수 및 깊이별로 독립적인 block-diagonal 구조를 가져야 함을 Schur's lemma를 이용해 엄밀히 증명(Theorem 3.1).
3. CRT 기반 채널 관련성 예측 및 실증: square-free 합성 모듈러스 N에 대해 CRT가 예측하는 task-relevant 소수 채널이 실제로 ablation 시 성능(0.60–0.92)에 큰 영향을 미치는 반면, task-irrelevant 채널은 통계적 잡음 수준 이하로 영향이 거의 없음을 소수 개수, 합성 모듈러스, 입력 범위에 대한 체계적 sweep을 통해 확인.

How

Figure 4. Experiment 1. Mean diagonal drop (ablate own prime).

• 소수 집합 P0(3부터 59까지 16개 소수, 2는 sin 항이 항상 0이 되어 제외)와 깊이 상한 D=6을 고정하여 총 임베딩 차원 192(소수당 24차원)로 설계
• 각 (p, d) 블록을 Z/p^{d+1}Z 상의 문자 χ1(a)=e^{2πia/p^{d+1}}의 실수 형태(2×2 회전 행렬로 작용)로 정의하고, 이를 addition이 회전으로 작용하는 equivariant 구조로 구성
• Schur's lemma(Appendix A, Theorem A.16)를 적용하여, 서로 다른 소수 또는 서로 다른 깊이에 대응하는 characters가 비동형(non-isomorphic) irreducible representation임을 보이고 이에 대한 equivariant intertwiner가 0이 됨을 증명(Theorem 3.1)
• Experiment 1에서 소수 채널별 diagonal drop(자기 프라임 ablation)과 off-diagonal drop(타 프라임 ablation), 그리고 수렴 속도(test acc > 0.85 도달)를 소수 개수·합성 모듈러스·입력 범위를 체계적으로 바꿔가며 측정

Originality

• 기존 FoNE(base-10 기반) 또는 Stevens et al.(2024)의 단일 주파수 기반 각도 임베딩과 달리, adelic Fourier 분해에 기반해 소수별로 독립적인 주파수 채널을 구성하는 원리적(principled) 접근을 제시
• 임베딩의 구조적 성질(block-diagonal)을 단순히 경험적으로 관찰하는 데 그치지 않고 Schur's lemma를 통한 대수적 증명으로 formalize
• CRT를 임베딩 설계와 해석의 예측 도구로 직접 연결하여, 어떤 프라임 채널이 특정 모듈러스 task에 관련되는지 사전에 예측 가능하게 함

Limitation & Further Study

• 실험에 사용된 소수 집합 P0(3~59)에 포함되지 않는 소수를 인수로 갖는 모듈러스는 모델이 원천적으로 인식할 수 없어, 임베딩의 확장성과 임의의 큰 소수·모듈러스에 대한 일반화 가능성이 제한적임
• square-free 합성 모듈러스에 초점을 맞추고 있어, 소수의 거듭제곱 인수(non-squarefree modulus)를 갖는 경우에 대한 이론적·실험적 검증이 부족함
• 논문 발췌 부분에서는 실제 신경망 아키텍처(분류기 등)와의 통합, 학습 성능이 실제 down-stream task(예: 산술 추론, 언어모델 pretraining)에서 어떻게 전이되는지에 대한 광범위한 실증이 제한적으로 제시되어 후속 연구가 필요함
• off-diagonal drop이 "대부분의 설정에서" 통계적 잡음 수준 이하라고 언급되어 있어, 일부 설정에서는 완전한 독립성이 보장되지 않을 가능성에 대한 추가 분석이 필요함

Evaluation

총평: 모듈러 산술 학습을 위한 임베딩 설계에 있어 조화해석과 표현론(Schur's lemma, CRT)을 활용해 이론적으로 근거 있는 구조를 제시하고 이를 체계적 실험으로 검증한 점에서 견고하고 참신한 연구이나, 소수 집합의 제한과 non-squarefree 모듈러스에 대한 확장성은 향후 과제로 남아있다.