How do Small Transformer Models Learn Hard Math Tasks?

Motivation

• Known: SALSA, PICANTE, VERDE, FRESCA 등의 선행 연구는 transformer가 LWE 샘플로부터 secret vector를 복구할 수 있음을 보였으며, 초기 encoder-decoder 모델에서는 cross-attention 가중치를 분석해 secret bit를 직접 읽어낼 수 있었다.
• Gap: 그러나 architecture가 효율적인 encoder-only 모델로 진화하면서 cross-attention 메커니즘이 사라졌고, secret이 정확히 어디에 어떻게 인코딩되는지에 대한 mechanistic한 이해가 부재했다.
• Why: 이 문제는 post-quantum cryptography 스킴의 보안 근거가 되는 LWE의 취약성을 이해하는 데 중요할 뿐 아니라, transformer가 표준 예측 경로가 실패할 때 어떻게 대안적 representational capacity를 할당하는지에 대한 근본적 통찰을 제공하기 때문에 중요하다.
• Approach: encoder-only transformer의 layer-wise 및 embedding-level mechanistic interpretability 분석을 dimensionality reduction, causal intervention, linear probing을 통해 최초로 수행하고, 이를 바탕으로 positional embedding에 L1 sparsity regularization을 적용하는 architectural intervention을 제안한다.

Achievement

1. LWE에 대한 최초의 mechanistic 분석: cryptanalysis task에 학습된 encoder-only transformer에 대해 layer-wise 및 embedding-level 수준의 최초의 mechanistic interpretability 분석을 제시했다.
2. Positional Memorization 현상 발견: 고노이즈 알고리즘 task 하에서 transformer가 latent state variable(secret)을 positional embedding 내부에 직접 memorize한다는 것을 밝혔다. 모델은 training objective(b 예측)에서는 near-zero exact accuracy를 보이지만 secret 복구는 성공한다.
3. Sparse Positional Regularization 제안: positional embedding에 L1 sparsity regularization을 적용하는 architectural modification을 통해, 계산 비용이 큰 post-hoc distinguisher 기반 secret 복구 과정을 human-interpretable한 파라미터 직접 검사로 대체했다.

How

• BKZ와 Flatter lattice reduction으로 전처리된 TAPAS LWE 데이터셋(N=256, log2 q=20)을 사용
• Hamming weight h∈{20,21,22}인 30개의 binary secret에 대해 encoder-only transformer를 학습시키고 예측 정확도(near-zero)를 확인
• Layer-wise PCA를 통해 residual stream 및 학습된 positional embedding 벡터의 구조를 시각화 (Figure 2, 3)
• Causal intervention(ablation)을 통해 secret position을 제거했을 때 성능에 미치는 영향을 측정 (Figure 4)
• Linear probing을 통해 positional embedding에 secret 정보가 선형적으로 인코딩되어 있는지 검증
• Positional embedding에 L1 sparsity regularization(α)을 적용하여 secret을 명시적으로 isolate하는 architectural intervention 도입 (Figure 5)

Originality

• 기존 SALSA/PICANTE/VERDE/FRESCA 연구가 post-hoc distinguisher를 통한 secret 복구에 집중했다면, 본 논문은 encoder-only 모델 내부의 representational 메커니즘 자체를 최초로 해부(mechanistic interpretability)한다는 점에서 독창적이다.
• 모델이 training objective를 사실상 풀지 못함에도 secret을 복구하는 역설적 현상을 발견하고, 이를 positional embedding에 대한 "hidden pathway"로 설명한 것은 새로운 관점이다.
• 해석 가능성 통찰을 단순 분석에 그치지 않고 L1 sparsity regularization이라는 architectural intervention으로 전환하여 실용적 해법(direct parameter inspection)으로 발전시킨 점이 독창적이다.

Limitation & Further Study

• 분석이 특정 파라미터 설정(N=256, log2 q=20)과 제한된 30개 secret 샘플에 국한되어 있어, 다른 LWE 파라미터나 더 큰 규모의 모델/차원에서도 동일한 positional memorization 현상이 일반화되는지 추가 검증이 필요하다.
• 본문 발췌만으로는 L1 regularization 도입이 기존 distinguisher 기반 방법 대비 얼마나 계산 비용을 절감하고 정확도를 유지하는지에 대한 정량적 비교가 충분히 제시되지 않았다.
• positional embedding에 secret이 저장되는 근본적 원인(왜 모델이 이런 경로를 선택하는지에 대한 이론적 설명)이 아직 부족하며, 다른 task(암호학 외 범용 알고리즘 task)로의 일반화 가능성에 대한 논의가 추가로 필요하다.

Evaluation

총평: 암호학적 하드 문제에 대한 transformer의 내부 동작을 최초로 mechanistic하게 규명하고, 이를 실용적인 architectural intervention으로 연결한 점에서 흥미롭고 의미 있는 연구이나, 워크샵 페이퍼 수준의 제한된 실험 범위와 일반화 검증 부족이 아쉽다.