Learning Large-Scale Modular Addition with an Auxiliary Modulus

Motivation

• Known: parity function 및 modular addition 학습은 input sensitivity 때문에 어려우며, N과 q가 커질수록 학습에 필요한 샘플 수와 network width가 증가한다는 것이 알려져 있다. 최근 Saxena et al.(2025)의 sparse method는 training sequence에 zero를 늘려 유효 summand 수를 줄임으로써 N, q를 크게 확장하는 데 성공했다.
• Gap: sparse method는 training 분포(sparse, 많은 zero)와 test 분포(uniform) 간 covariate shift를 유발하며, 이로 인한 일반화 오차가 이론적으로 규명되지 않았고 Dropout, PreNorm, Bias, weight initialization 등 모델 설정에도 민감하게 반응하는 문제가 있다.
• Why: modular addition 학습은 post-quantum cryptography 기반인 Learning With Errors(LWE) 문제 공격의 핵심 요소로 지목되었으며, large modulus q에 대한 학습 능력 확장이 이 응용에서 매우 중요하기 때문에 covariate shift 없이 확장 가능한 학습법은 실용적 중요성이 크다.
• Approach: training 시 원래 modulus q 외에 K>1배 확대된 auxiliary modulus Kq에 대한 modular addition을 확률 r로 main loss 대신 학습시켜, 입력 분포는 그대로 유지하면서 wrap-around 빈도와 문제 난이도만 낮추는 방법을 제안한다.

Achievement

Figure 4. Heatmaps of match accuracy across various combinations of K and r for token embedding. The color scale represe

1. 이론적 분석: sparse method의 covariate shift로 인한 일반화 오차 하한을 Proposition 1로 도출하고, (1-1/q)^N이 일반적인 modular addition 환경에서 크게 나타남을 보였다(Figure 1a).
2. 모델 설정 민감성 규명: sparse method가 Dropout, PreNorm, Bias, weight initialization 등 구성 요소 변화에 취약함을 실험적으로 확인하였다(Figure 1b).
3. 새로운 covariate-shift-free 방법 제안: auxiliary modulus Kq를 도입한 loss를 확률 r로 main loss와 교체하여 학습 분포를 유지하면서 난이도를 조절하는 기법을 제시하였다.
4. 대규모 스케일링 및 샘플 효율성 입증: N=64, q=974269 조건에서 100K 샘플만으로 97.0%의 τ-accuracy(τ=0.05)를 달성했으며, 이는 동일 데이터 규모에서 sparse method의 9.5%, 그리고 1M 샘플로 확장했을 때의 93.9%보다도 우수한 결과이다.

How

Figure 4. Heatmaps of match accuracy across various combinations of K and r for token embedding. The color scale represe

• Modular addition을 fq(x) = Σxi mod q로 정의하고, 기존 sparse method(Saxena et al., 2025)의 sparse sampling 절차(z ~ g(z) ∝ 1/√(N-z+1))를 분석
• Proposition 1을 통해 zero-free input에서의 error가 test-time generalization gap으로 전이됨을 증명(증명 스케치 제공, 전체 증명은 부록 Section A)
• Dropout, PreNorm, Bias, weight initialization 등 다양한 model configuration에서 sparse method 성능을 비교하는 실험 수행(Figure 1b)
• 제안 방법에서 auxiliary modulus Kq에 대한 loss를 확률 r로 main loss와 교체하는 학습 절차 설계
• K와 r의 다양한 조합에 대해 match accuracy를 측정하는 heatmap 실험 수행(Figure 3, 4, 5)
• ρ 지표에 대한 heatmap 분석을 통해 방법의 특성을 평가(Figure 2)
• N=64, q=974269 등 대규모 설정에서 학습 데이터 크기(100K~1M)에 따른 match accuracy 및 τ-accuracy 비교

Originality

• 기존 sparse method의 covariate shift 문제를 최초로 이론적(Proposition 1)·실증적으로 규명
• 입력 분포를 변경하지 않으면서 문제 난이도를 조절하는 auxiliary modulus Kq라는 새로운 아이디어 제시
• 기존 curriculum learning이나 sparse input 기반 방법과 달리, 동일 문제의 다른 modulus를 auxiliary task로 활용하는 접근은 관련 연구(Both et al., 2026 등)와 차별화됨
• angle embedding 등 입력 자체를 바꾸는 방법과 달리 auxiliary task를 통해 periodicity 학습을 유도하는 방식으로 기존 연구와 구별됨

Limitation & Further Study

• auxiliary modulus K와 replace probability r을 어떻게 최적으로 선택할지에 대한 이론적 가이드라인이 충분히 제시되지 않고 heatmap 기반 경험적 탐색에 의존하는 것으로 보임
• 실험이 특정 N, q 조합(예: N=64, q=974269)에 집중되어 있어 더 다양한 실제 응용(LWE 공격 등)으로의 일반화 가능성에 대한 추가 검증이 필요함
• 발췌된 abstract 및 본문만으로는 모델 아키텍처의 세부 사항과 전체 실험 조건(예: transformer 구조, 하이퍼파라미터)이 명확히 파악되지 않음
• auxiliary modulus 방법이 q가 소수가 아니거나 다른 대수적 구조를 가질 때도 잘 작동하는지에 대한 추가 분석이 필요해 보임

Evaluation

총평: covariate shift라는 기존 방법의 숨겨진 문제를 이론과 실험으로 명확히 규명하고, 이를 해결하는 간단하면서도 효과적인 auxiliary modulus 기법을 제안하여 대규모 modular addition 학습의 샘플 효율성과 확장성을 크게 개선한 견실한 연구이다.