경험적 NTK로 특징 찾기

초록

본 논문은 학습된 신경망의 경험적 신경탄젠트 커널(eNTK) 고유벡터가 모델이 실제로 사용하는 특징과 일치한다는 증거를 제시한다. TMS, 1‑층 MLP 및 1‑층 Transformer의 세 가지 표준 토이 모델에서 eNTK의 상위 고유공간이 실제 특징(특히 Fourier 특징)과 정렬되는 것을 발견했으며, 층별 eNTK를 통해 특징이 어느 층에 존재하는지, 그리고 grokking 전이와 같은 학습 단계 변화를 감지할 수 있음을 보였다.

상세 분석

논문은 경험적 신경탄젠트 커널(eNTK)이 학습 종료 시점에 모델의 함수 근사에 여전히 유효하다는 ‘eNTK 가설’을 전제로, 고유벡터가 의미 있는 특징을 담고 있는지를 실험적으로 검증한다. 첫 번째 실험인 Toy Models of Superposition(TMS)에서는 입력‑출력 차원이 n=50인 1‑층 오토인코더를 다양한 은닉 차원 m과 희소도 S에서 학습시켰다. eNTK의 평탄화된 스펙트럼은 두 개의 뚜렷한 ‘클리프’를 보였으며, 두 번째 클리프가 정확히 n=50번째 고유값에서 발생한다. 이는 고유공간 차원이 데이터에 존재하는 실제 특징 수와 일치함을 의미한다. 고유벡터와 각 특징의 활성화 벡터 사이의 코사인 유사도를 히트맵으로 시각화했을 때, β=1(중요도 가중)에서는 고중요도 특징이, β=0.3에서는 저중요도 특징까지도 거의 일대일 대응을 이루는 것을 확인했다. 이는 eNTK가 원본 커널 자체에 모든 특징 정보를 보유하고 있으나, 손실 가중치에 따라 강조 정도가 달라진다는 점을 시사한다.

두 번째 실험에서는 1‑층 MLP가 모듈러 덧셈(mod p) 문제를 학습하도록 설정하였다. p=29, 은닉 유닛 n=512, 학습 비율 α=0.7에서 grokking 현상이 관찰되었다. eNTK 스펙트럼은 두 개의 급격한 하강(클리프)을 보였는데, 첫 번째 클리프는 차원 k=4⌊p/2⌋≈56으로, 이는 첫 번째 층이 학습한 Fourier 기저(코사인·사인 조합)와 일치한다. 두 번째 클리프는 학습 후반, 테스트 정확도가 급상승하는 grokking 전이 시점에 나타났으며, 이는 두 번째 층이 ‘합’과 ‘차’ Fourier 특징을 학습함을 의미한다. 층별 eNTK를 계산하면 첫 번째 클리프는 입력‑첫 번째 은닉 층에, 두 번째 클리프는 두 번째 은닉 층에 국한되어 나타나, 특징이 층별로 명확히 구분됨을 보여준다. 이는 eNTK 스펙트럼이 모델 내부 표현의 구조적 변화를 진단하는 도구로 활용될 수 있음을 뒷받침한다.

세 번째 실험은 1‑층 Transformer가 동일한 모듈러 덧셈을 수행하도록 학습시킨 경우이다. 여러 시드에 대해 실험했으며, eNTK의 주요 고유벡터는 Attention 블록의 O/V 레이어, MLP 블록의 입력·출력 레이어, 그리고 언임베딩 레이어에 국한되었다. 각 고유벡터는 모델이 선택한 ‘키 주파수’(Fourier 모드)와 높은 정렬도를 보였으며, 이는 Transformer가 복잡한 연산을 수행하면서도 여전히 Fourier 기반의 선형 변환을 핵심 특징으로 활용한다는 점을 시사한다. 특히, eNTK가 층별로 분해될 때 각 층이 담당하는 주파수 집합이 서로 겹치지 않으며, 이는 모델이 기능을 모듈화하여 학습한다는 가설을 뒷받침한다.

전반적으로 논문은 eNTK 고유벡터가 단순히 수학적 도구가 아니라, 실제 학습된 특징을 드러내는 ‘특징 탐지기’로 작동한다는 강력한 증거를 제공한다. 스펙트럼의 클리프 구조는 특징 수와 직접 연결되며, 층별 eNTK는 특징의 위치와 학습 단계 변화를 동시에 파악한다. 이러한 결과는 작은 모델에 대한 메커니즘 해석뿐 아니라, 대규모 모델에서도 커널 기반 진단 기법을 확장할 가능성을 열어준다. 다만, 현재 실험은 제한된 토이 모델에 국한되어 있어, 실제 자연어 처리나 비전 모델에 적용하기 위한 추가 연구가 필요하다.