A Counterexample to Fourier Alignment in Single-Neuron Modular Addition

기술 요약: 단일 뉴런 모듈러 덧셈에서의 푸리에 정렬에 대한 반례

문제 정의
본 논문은 모듈러 산술(modular arithmetic)을 학습하는 신경망의 표현론적 특성에 관한 구체적인 질문을 다룹니다. 이 과제는 '그로킹(grokking)'과 알고리즘 일반화의 핵심인 순환군 $\mathbb{Z}/p\mathbb{Z}$ (여기서 $p \ge 5$는 소수) 상의 덧셈에 대해, 정확한 전체 테이블 교차 엔트로피(exact full-table cross-entropy)로 학습된 단일 ReLU 뉴런에 관한 문제를 조사합니다. 구체적으로는 [Cla26d]에서 제기된 문제, 즉 뉴런이 초기에 활성화되어 있다는 조건 하에 정규화된 파라미터 벡터 $(u, v, w)$가 반드시 단일 비제로 실수 푸리에 클래스(즉, 특정 주파수 $\zeta$와 그 켤레 $-\zeta$)에 의해 지지되는 방향으로 수렴해야 하는지에 대해 조사합니다. 이전의 2층 네트워크에 대한 이론적 연구들은 개별 뉴런이 단일 주파수 푸리에 특징을 선택하는 경향이 있으며, 주파수들이 초기 스펙트럼 크기에 따라 경쟁한다고 제안했습니다. 본 논문은 가장 단순한 모델인 단일 뉴런에 대해 이러한 "단일 주파수 정렬"이 훈련 역학의 필연적인 결과인지 여부를 결정하고자 합니다.

방법론
저자는 ReLU 활성화 함수의 0에서의 비미분성을 인정하며, 연속 시간 훈련 역학을 모델링하기 위해 클라크 미분 포함(Clarke differential inclusions)을 사용하는 엄격한 분석적 접근 방을 채택합니다. 방법론은 다음과 같습니다:

1. 특정 초기 조건의 구축: 저자는 정확히 하나의 입력 게이트(입력 $(0,0)$에 대응)가 활성화되어 있고, 나머지 모든 게이트는 엄격하게 비활성 상태(음의 사전 활성화)인 초기 파라미터 집합을 정의합니다.
2. 동역학계 분석: 저자는 크로스 엔트로피 손실의 그래디언트 흐름(gradient flow)을 분석합니다. 대칭성을 활용하여, 고차원 시스템을 활성 게이트의 사전 활성화 및 출력 가중치를 제어하는 저차원 상미분 방정식(ODE)으로 축소합니다.
3. 불변량 분석: 저자는 출력 가중치 차이와 입력 사전 활성화 사이의 관계를 나타내는 보존량(또는 특정 단조성을 가진 양)을 식별합니다. 이를 통해 활성 게이트의 사전 활성화가 유한한 시간 내에 0으로 선형적으로 감소함을 증명할 수 있습니다.
4. 클라크 궤적 선택: 사전 활성화가 0에 도달하면, 저자는 클라크 일반화 미분 프레임워크를 사용합니다. 저자는 0에서의 ReLU 미분을 0으로 선택함으로써(자동 미분에서 표준적인 관례), 궤적이 모든 그래디언트가 사라지는 "죽은(dead)" 상태에 진입하여 파라미터가 고정됨을 보여줍니다.
5. 푸리에 스펙트럼 분석: 저자는 고정된 파라미터 벡터의 푸리에 에너지 분포를 계산합니다. 이들은 극한의 방향이 중심 점 질량 벡터(centered point mass vector)의 스칼라 배이며, 이는 모든 비제로 실수 푸리에 클래스에 에너지를 균등하게 분배함을 보여줍니다.

주요 기여 및 결과

• 단일 주파수 정렬의 반박: 본 논문은 [Cla26d]에서 제기된 문제에 대한 부정적인 해답을 제공합니다. 저자는 정규화된 파라미터가 단일 푸리에 클래스에 의해 지지되지 않는 극한으로 수렴하는 반례를 구축합니다. 대신, 푸리에 에너지는 모든 $p-1$개의 비제로 주파수 클래스에 균등하게 분배됩니다. $p \ge 5$인 경우, 단일 클래스에 포함된 에너지의 비율은 $2/(p-1)$이며, 이는 1보다 엄격히 작습니다.
• 유한 시간 사망 및 고정(Finite-Time Death and Freezing): 이 반례는 단일 뉴런이 유한한 시간 내에 완전히 비활성 상태("죽은 ReLU")에 도달할 수 있음을 보여줍니다. 일단 죽게 되면, 파라미터는 초기 조건과 죽음에 이르는 역학에 의해 결정된 극한값에서 고정됩니다.
• 강건성 및 확률: 이 반례는 고립된 인위적 결과가 아닙니다.
  • 이 결과는 초기 조건의 열린 집합(open set) 상에서 성립하며, 이는 표준 가우시안 초기화 하에서 양의 확률로 발생함을 의미합니다.
  • 부록 A는 다음을 보여줌으로써 결과를 강화합니다:
    • 모든 클라크 궤적(특정 궤적 선택의 필요성 제거)에 대해 열린 집합의 초기 조건에서 실패가 발생합니다.
    • 이 현상은 $\text{ReLU}'(0)=0$인 관례 하에서도 지속됩니다.
    • 이 현상은 ReLU의 매끄러운 데드 존 근사 함수($z \le 0$에서 0인 $C^\infty$ 함수)에서도 유지됩니다.
    • 이는 뉴런이 단 한 단계의 스텝에서 죽는 고정 스텝 풀 배치 그래디언트 디센트(fixed-step full-batch gradient descent)에서도 발생합니다.
• 지속적 활성 반례: 부록은 뉴런이 활성 상태를 유지하지만(하나의 게이트가 양수로 유지됨) 정규화된 가중치가 여전히 다중 주파수 극한으로 수렴하여, 뉴런이 죽지 않고도 단일 주파수 정렬에 실패하는 "지속적 슬라이딩(persistent sliding)" 궤적을 제시합니다.

의의 및 주장
본 논문은 단일 주파수 또는 모노세만틱(monosemantic) 푸리에 구조가 모듈러 덧셈을 학습하는 단일 뉴런의 일반적인 결과가 아님을 주장합니다.

• 실패 메커니즘: 실패는 두 가지 익숙한 메커니즘에서 기인합니다: 뉴런이 죽거나(데드 영역 진입), 또는 뉴런이 활성 상태를 유지하면서 일반적인 알고리즘을 배우는 대신 단 하나의 테이블 엔트리(이 경우 $(0,0)$ 케이스)를 "암기"하는 것입니다. 지속적 궤적은 가장 작은 모델에서 관찰되는 암기-일반화 이분법 중 "암기" 측면을 실현합니다.
• 한계: 저자는 자신의 결과가 활성화 함수의 명확한 데드 영역 존재에 의존하고 있음을 명시적으로 언급합니다. 모든 곳에서 양수인 활성화 함수(예: softplus)에 대한 문제는 여전히 미해결 과제로 남아 있습니다.
• 맥락: 이 연구는 푸리에 구조가 해 공간(solution space)에 존재함에도 불구하고, 훈련 역학이 단일 주파수 선택을 강제하지는 않는다는 점을 명확히 합니다. 결과적으로 단일 뉴런 모델과 더 넓은 네트워크(예: [MEO+24]의 너비 $m \ge 4(p-1)$)에서 관찰되는 다중 주파수 선택 메커니즘을 구분합니다.

요약하자면, 본 논문은 단일 주 frequency 정렬이 필연적이라는 가설이 거짓임을 입증하며, 극한의 방향이 푸리에 클래스 측면에서 최대한으로 순수하지 않은 구체적인 반례를 제시합니다.