Neural Networks Provably Learn Spectral Representations for Group Composition

기술 요약: 신경망은 군 합성(Group Composition)을 위한 스펙트럼 표현을 증명 가능하게 학습한다

문제 정의
본 논문은 신경망의 훈련 과정 중 발생하는 구조화된 내부 표현의 출현을 조사하며, 특히 군 합성(group composition) 과업에 초점을 맞춘다. 유한 군 $(G, \star)$가 주어졌을 때, 목표는 임의의 쌍 $(g_1, g_2) \in G \times G$에 대해 곱 $g_1 \star g_2$를 예측하도록 하는 2층 신경망을 훈련하는 것이다. 이전 연구들은 순환 군($\mathbb{Z}_p$)에서의 모듈로 덧셈을 학습하는 신경망이 푸리에 특징(단일 주파수 특화)을 발견한다는 것을 확립했으나, 임의의 유한 군—특히 표현이 스칼라가 아닌 행렬 값인 비가환 군(non-Abelian groups)—을 학습하는 신경망의 동작은 이론적으로 규명되지 않은 상태였다. 핵심 질문은 일반적인 군 $G$의 기약 표현(irreducible representations, irreps)을 신경망이 실제로 발견하는지, 그리고 아벨 군(Abelian settings)에서 관찰된 위상 정렬(phase alignment) 및 저계수 구조(low-rank structure)의 원리가 이 더 넓은 맥락으로 어떻게 일반화되는가 하는 점이다.

방법론
저자들은 유한 군 $G$의 완전한 합성표(composition table)를 크로스 엔트로피 손실(cross-entropy loss)을 사용하여 학습하는, 이차 활성화 함수($\sigma(x) = x^2$)를 가진 2층 완전 연결 신경망을 분석한다. 신경망 아키텍처는 비가환 연산을 수용하기 위해 두 피연산자에 대한 별도의 입력 임베딩을 사용한다.

이론적 분석은 두 단계로 진행된다:

1. 단계 I (특징 학습): 뉴런의 스케일링 인자(scaling factors)를 작은 상수 $a$로 고정하고, 방향 파라미터(입력 및 출력 임베딩)를 단위 구면 상의 투영 경사 흐름(projected gradient flow)을 통해 진화시킨다. 저자들은 유클리드 파라미터 공간으로부터의 훈련 역학을 군 이산 푸리에 변환(Group DFT)을 사용하여 스펙트럼 영역으로 들어 올린다(lift). 이들은 투영 경사 흐름이 푸리에 계수의 매니폴드 상에 정의된 특정 에너지 범함수 $\Omega_m$에 대한 리만 경사 상승(Riemannian gradient ascent)과 동등함을 입증한다.
2. 단계 II (스케일 최적화): 방향 파라미터가 수렴하면, 마진(margin)을 최대화하여 크로스 엔트로피 손실을 0으로 몰아넣기 위해 스케일링 인자를 최적화한다.

분석은 유한 군에서의 조화 해석(harmonic analysis), 특히 기약 표현의 직교성 및 슈어의 직교 관계(Schur's orthogonality relations)에 크게 의존한다. 저자들은 동역학계(dynamical systems) 도구를 사용하여 임계점(평형점)의 분류와 연속 리만 경사 흐름으로 확장된 안장점 회피(saddle avoidance) 논증을 활용하여 수렴 특성을 증명한다.

주요 기여 및 결과

• 일반 군 표현 학습 (정리 4.3): 임의의 유한 군(비가환 군 포함)에 대하여, 저자들은 무작위 초기화 하에서 경사 흐름이 각 뉴런을 다음과 같은 상태로 거의 확실하게(almost surely) 수렴하도록 유도함을 증명한다:

  • 단일 표현: 각 뉴런은 정확히 하나의 비자명 기약 표현 $\check{\rho}_m$(및 그 쌍대 $\check{\rho}_m^\vee$)을 선택하며, 이로 인해 다른 표현에 대응하는 모든 푸리에 계수는 소멸한다.
  • 계수 1의 회전 정렬(Rank-One Rotational Alignment): 생존한 표현의 푸리에 계수는 계수 1(rank-one) 행렬이 된다. 또한, 이 계수들은 입력 및 출력 층 전반에 걸쳐 특정 회전 정렬을 보인다: $\hat{\xi}_m[\rho] \propto \hat{\theta}_{2m}[\rho] \hat{\theta}_{1m}[\rho]$ 및 이들의 순환 치환. 이는 모듈로 덧셈에서 관찰된 스칼라 위상 정렬을 행렬 값 표현으로 일반화한 것이다.
  • 메커니즘: 증명은 훈련 역학이 에너지 범함수를 최대화함을 확립한다. 비대상 평형점(음의 에너지, 자명한 지지도를 가진 제로 에너지, 또는 높은 계수의 양의 에너지)은 측도 0의 트랩(measure-zero traps)이거나 엄격한 안장점(strict saddles)임이 밝혀졌으며, 경사 흐름에 의해 거의 확실하게 회피된다.

• 아벨 군 특성 (정리 5.1 및 5.3): 아벨 군에 대해, 저자들은 모집단 수준의 완전한 설명을 제공한다:

  • 다양화: 균등 구면 초기화 하에서, 학습된 표현들은 모든 비자명 기약 표현 상에 균등하게 분포하며, 위상은 선택된 표현과 독립적으로 단위 원 상에서 하르 균등(Haar-uniform) 분포를 따른다.
  • 결함 있는 지표(Flawed Indicator): 결과적인 앙상블 예측기는 올바른 군 곱 $g_1 \star g_2$에 가장 높은 로짓(logit)을 할당하며, $g_1^2$ 및 $g_2^2$에서 보조적인 "유령" 피크(ghost peaks)를 갖는 "결함 있는 지표" 역할을 한다. 이러한 구조는 표현의 직교성과 하르 균등 위상의 평균화를 통한 노이즈 항의 정확한 상쇄로부터 발생한다.
  • 수렴 속도: 논문은 위상 정렬 및 표현 경쟁에 대한 명시적인 지수적 수렴 속도를 확립한다. 승리하는 표현의 선택은 초기 스펙트럼 크기에 의해 결정되는 "로터리 티켓(lottery ticket)" 메커니즘을 따른다.

• 단계 II 역학 (정리 4.5): 저자들은 평균장 예측기(mean-field predictor)가 완벽한 정확도 조건을 만족하면(이는 아벨 군에서 성립하며 비가환 군에서도 경험적으로 검증됨), 스케일링 인자가 로그 함수적으로 성장하여 소프트맥스 분포를 날카롭게 만들고 $O(1/T)$의 속도로 손실을 0으로 몰아넣는다고 증명한다.

의의 및 주장
본 논문은 일반적인 유한 군에 대해 신경망이 군 합성을 어떻게 학습하는지에 대한 최초의 엄밀한 이론적 규명을 제공한다고 주장한다. 순환 군에서의 "단일 주파수 특화" 현상을 더 넓은 기약 표현 프레임워크로 통합함으로써, 저자들은 "푸리에 특징"이 모듈로 덧셈에서의 스칼라 케이스가 단순히 기약 표현을 학습하는 것임을 보여준다.

그 의의는 다음과 같다:

1. 기계적 해석 가능성과 표현론의 가교: 본 연구는 창발된 신경망 특징과 기저 군의 수학적 구조 사이의 연결을 공식화하여, 모듈로 덧합에서의 "푸리에 특징"이 기약 표현을 학습하는 것의 단순한 스칼라 사례임을 보여준다.
2. 새로운 저계수 압축: 비가환 군의 행렬 값 표현에서 계수 1 회전 정렬을 발견한 것은, 다른 딥러닝 맥락에서 관찰된 저계수 편향(low-rank bias)을 스펙트럼 영역으로 확장한 새로운 구조적 발견이다.
3. 역학적 설명: 이 논문은 단순한 경험적 관찰을 넘어, 무작위 초기화로부터 왜, 그리고 어떻게 이러한 구조화된 솔루션이 나타나는지를 리만 경사 흐름 및 안장점 회피를 통한 동역학계 기반의 증명을 통해 설명한다.

저자들은 일반적인 유한 군을 다루지만, 고차원 기약 표현이나 자기 공액(self-conjugate) 표현을 가진 군에 대한 한계 분포의 이론적 규명, 그리고 유한 샘플 설정에서의 일반화 간극(grokking) 분석에 관한 미해결 과제가 남아 있음을 언급한다.