[논문 리뷰] Towards Understanding Grokking: An Effective Theory of Representation Learning

Ziming Liu, O. Kitouni외 4명|arXiv (Cornell University)|2022. 05. 20.

Neural dynamics and brain function인용 수 24

한 줄 요약

본 논문은 grokking을 설명하기 위해 효과 이론과 상태도(phase diagrams)를 개발하고, 일반화와 구조화된 표현 간의 연관성을 제시하며 트랜스포머 및 토이 모델에서 네 가지 학습 단계(이해, grokking, 기억화, 혼란)를 식별한다.

ABSTRACT

We aim to understand grokking, a phenomenon where models generalize long after overfitting their training set. We present both a microscopic analysis anchored by an effective theory and a macroscopic analysis of phase diagrams describing learning performance across hyperparameters. We find that generalization originates from structured representations whose training dynamics and dependence on training set size can be predicted by our effective theory in a toy setting. We observe empirically the presence of four learning phases: comprehension, grokking, memorization, and confusion. We find representation learning to occur only in a "Goldilocks zone" (including comprehension and grokking) between memorization and confusion. We find on transformers the grokking phase stays closer to the memorization phase (compared to the comprehension phase), leading to delayed generalization. The Goldilocks phase is reminiscent of "intelligence from starvation" in Darwinian evolution, where resource limitations drive discovery of more efficient solutions. This study not only provides intuitive explanations of the origin of grokking, but also highlights the usefulness of physics-inspired tools, e.g., effective theories and phase diagrams, for understanding deep learning.

연구 동기 및 목표

• 작은 알고리즘 데이터세트에서 과적합한 후 신경망이 왜 일반화하는지 조사한다.
• 일반화를 구조화된 입력 표현의 등장과 연결한다.
• 임계 학습 데이터 크기와 학습 궤적을 예측하기 위한 물리학에서 영감을 받은 효과 이론을 개발한다.
• 학습 단계와 하이퍼파라미터가 지연 일반화에 미치는 영향을 특성화한다.

제안 방법

• 입력이 학습 가능한 임베딩에 매핑되고 디코더가 이들의 합을 처리하는 단순화된 토이 모델을 구성한다.
• 임베딩의 평행사변형 구조를 기반으로 Representation Quality Index (RQI)을 정의한다.
• 임베딩 동역학을 지배하는 효과적 손실 ell_eff와 Hessian 스펙트럼으로부터 grokking 속도를 도출하는 것을 공식화한다.
• 토이 및 트랜스포머 유사 설정에서 표현 학습 속도와 디코더 학습 속도 및 가중치 감소를 변화시키며 상태도를 도출한다.
• 작업 간 일반성(test)을 확인하기 위해 모듈러 덧셈 및 비아벨 그룹으로 분석을 확장한다.

실험 결과

연구 질문

• RQ1 모델이 학습 데이터에 과적합될 때 grokking에서 일반화의 원인은 무엇인가?
• RQ2 학습 데이터 크기가 구조화된 표현의 등장과 일반화에 어떻게 영향을 미치는가?
• RQ3 지연 일반화가 어떤 조건에서 발생하며 하이퍼파라미터가 이를 지연시키거나 가속시킬 수 있는가?
• RQ4 상태도와 효과 이론의 예측이 토이 모델을 넘어서 트랜스포머 아키텍처와 MNIST에서도 성립하는가?
• RQ5 다양한 연산에서 일반화를 가능하게 하는 표현 구조의 역할은 무엇인가?

주요 결과

• 일반화는 표현 공간에서 평행사변형을 만드는 구조화된 임베딩의 등장과 상관관계가 있다.
• 임계 학습 데이터 비율은 선형적이고 고유한 표현이 언제 나타나는지 결정하여 일반화를 가능하게 한다.
• Grokking은 이해와 기억화 사이의 단계이며, 적절한 하이퍼파라미터 튜닝으로 이동시키거나 제거할 수 있다.
• 상태도는 네 가지 학습 단계를 드러내며, 더 빠른 디코더나 과도한 디코더 용량이 기억화를 촉진하고, 균형 잡힌 학습이 이해와 grokking을 지지한다를 보여준다.
• 트랜스포머 설정에서 디코더 용량을 줄이거나 정규화를 적용하면 일반화가 빨라지고 grokking을 억제할 수 있으며, MNIST 실험에서 그 근거가 제시된다.
• 효과 이론은 grokking 시간과 상전이 지점을 예측하며 신경망 학습 궤적과 질적으로 일치한다.