기술 요약: 그로킹(Grokking)은 가치가 있는가?

문제 제기

대규모 언어 모델(LLM)은 종종 "두 단계 추론의 저주(curse of two-hop reasoning)"를 보입니다. 즉, 개별적으로 학습된 사실들을 일관된 추론 사슬로 결합하는 데 실패하는 현상입니다. 예를 들어, 모델이 "배럭의 아내는 미셸이다"와 "미셸은 1964년에 태어났다"는 사실은 기억하지만, 내부 연산을 통해 "배럭의 아내는 1964년에 태어났다"는 사실을 추론해내지 못하는 경우입니다. 이러한 실패는 표현의 불일치(representational mismatch)에 기인합니다. 표준 트랜스포머는 얕은 층에서 원자적 사실(atomic facts)을 검색하지만, 더 깊은 층은 이러한 특정 잠재 표현을 후속 검색 단계를 위한 입력으로 활용하도록 훈련되지 않았습니다.

최근 문헌은 파라미터 공유 트랜스포머(parameter-sharing Transformer)가 "그로킹"이라 불리는 장기적인 훈련 단계를 거침으로써 이를 극복할 수 있다고 제안합니다. 이 과정 동안 "일반화 회로(Generalization Circuit)"가 출현합니다. 이 회로는 모델이 중간 층에서 중간(가교) 엔티티를 명시적으로 복구함으로써 추론 사슬을 내재화할 수 있게 합니다. 그러나 근본적인 의문이 남아 있습니다: 그로킹된 모델이 정말로 비-그로킹된 모델보다 다운스트림 작업에서 우월한가? 그로킹 단계를 기다리기 위한 막대한 계산 비용은 정당화될 수 있는가? 그로킹을 통해 습득된 추론 능력이 새로운 사실로 전이되는가?

방법론

저자들은 합성 조성 추론(synthetic compositional reasoning) 과업에 대해 최첨단 파라미터 공유 트랜스포머 아키텍처를 사용하여 기계론적 연구를 수행합니다.

• 데이터셋 구축: 2,000개의 엔티티와 200개의 고유 관계를 포함하여 총 40,000개의 원자적 사실로 구성된 합성 데이터셋을 구축했습니다. 데이터는 분포 내(In-Distribution, ID) 및 분포 외(Out-of-Distribution, OOD) 서브셋으로 분할됩니다. 모델은 모든 원자적 쿼리에 대해 훈련되지만, 조성적 쿼리는 오직 ID 사실들로만 구성된 것들로만 훈련됩니다. 평가는 진정한 일반화 여부(단순 암기가 아닌)를 테스트하기 위해 OOD 사실들이 포함된 조성적 쿼리에 대해 수행됩니다.
• 기계론적 분석: 연구는 모든 층의 은닉 상태를 디코딩하여 일반화 회로의 성숙을 시각화하는 "로짓 렌즈(logit lens)" 기법을 활용합니다. 추론 트레이스는 디코딩된 은닉 상태가 실제 정답 가교 엔티티와 일치할 경우 "가교 발견(Bridge found)"으로 분류됩니다.
• 실험적 변형:
  1. 자연적 그로킹(Natural Grokking): 일반적인 ID 전용 조성 데이터를 학습하여 그로킹의 자연스러운 출현을 관찰합니다.
  2. 전체 감독 증강(Full Supervision Augmentation): OOD 사실에 대한 조성적 감독을 제공하여(첫 번째 홉 전체 증강 및 두 번째 홉 전체 증강) 학습을 가속화하고 추론 경로가 변화하는지 테스트합니다.
  3. 전이성 평가(Transferability Evaluation): 사전 훈련된 모델들(자연적으로 그로킹된 모델과 "가짜 그로킹된" 모델 모두)을 유지된 ID 사실들과 2,000개의 완전히 새로운 원자적 사실들의 혼합 데이터로 미세 조정하여, 새로운 지식을 흡수하고 이를 조합하는 능력을 평가합니다.

주요 결과

1. 동일한 추론 경로

그로킹이 새로운 추론 패러다임의 갑작스러운 습득이라는 가설과는 반대로, 본 연구는 비-그로킹 모델과 그로킹 모델이 분포 내 조성 쿼리에 대해 동일한 추론 경로를 따른다는 것을 발견했습니다.

• "일반화 회로"는 그로킹 후에만 형성되는 새로운 구조가 아니라, ID 팩트에 대해 훈련 초기 단계에 이미 확립된 추론 경로입니다.
• 그로킹은 희소하게 감독된 OOD 사실들을 이 자연스럽게 확립된 회로에 통합하는 과정으로 정의됩니다.
• OOD 팩트에 대해 완전한 조성적 감독을 받은 모델은 훨씬 적은 단계만으로도 자연적으로 그로킹된 모델과 동일한 추론 경로 및 전이 성능을 달s합니다.

2. 행동과 메커니즘의 해리(Dissociation)

본 논문은 행동적 그로킹(보지 못한 사례에 대한 높은 정확도)이 일반화 회로의 형성과 본질적으로 연결되어 있지 않음을 입증합니다.

• 특정 데이터 환경(예: 두 번째 홉 전체 증강) 하에서, 모델은 중간 층에서 가교 엔티티를 올바르게 복구하지 못하면서도 OOD 조성에 대해 높은 정확도를 달성할 수 있습니다.
• 이러한 "가짜 그로킹된(fake grokked)" 모델들은 일반화 회로라는 기저 메커니즘이 결여되어 있으며, 결정적으로 실제 회로를 가진 모델과 달리 새로운 사실로 이 능력을 전이시키지 못합니다.

3. 제한된 전이성

성숙한 일반화 회로가 존재하더라도, 새로운 지식에 대한 전이성은 제한적입니다.

• 회로는 모델이 새로운 사실에 대해 가교 엔티티를 복구할 수 있게 해주지만, 두 번째 홉이 새로운 사실을 포함하는 경우 두 번째 홉의 추론을 완료할 수 있다는 것을 보장하지는 않습니다.
• 새로운 지식을 위한 신뢰할 수 있는 회로를 구축하기 위해서는, 모델이 여전히 추가적인 그로킹 단계나 해당 사실 및 그 조성을 다루는 실질적인 미세 조정 데이터가 필요합니다.
• "가짜 그로킹된" 모델은 새로운 사실과 기존 사실을 섞은 조성적 추론에 적응하지 못하고, 단순히 새로운 원자적 사실을 암기하는 데 그칩니다.

의의 및 주장

본 논문은 그로킹이 진정한 일반화된 추론 패러다임의 습득을 나타낸다는 통념에 도전합니다. 대신, 그로킹은 기존 구조에 희소한 데이터를 통합하는 메커니즘이라고 상정합니다.

저자들은 그로킹을 추구할지 여부에 대한 결정은 데이터 희소성과 계산 자원 사이의 트레이드오프를 포함한다고 결론짓습니다:

• 희소한 감독 하에서는: 사실을 추론 회로에 강제로 통합하기 위해 그로킹이 일어날 때까지 훈련하는 것이 필요할 수 있습니다.
• 충분한 감독 하에서는: 충분한 조성적 데이터가 있다면, 모델은 훈련 포화 상태에서 멈추고 긴 그로킹 단계의 계산 오버헤드 없이 동등한 성능을 달성할 수 있습니다.

궁극적으로, 본 연구는 "그로킹된" 트랜스포머가 인간과 같은 추론에 필적하는 완전한 조성 논리 숙달에 도달하지 못하며, 도메인을 가로질러 추론 패턴을 일반화하는 능력이 여전히 제약되어 있음을 시사합니다.

한계점

저자들은 두 가지 주요 한계를 인정합니다:

1. 계산적 제약: "가짜 그로킹"에 대한 조사는 그로킹 현상을 관찰하기 위해 필요한 극단적인 훈련 주기로 인해 철저히 이루어지지 못했으며, 이로 인해 회로 붕괴의 정확한 경계 조건은 미해결 과제로 남아 있습니다.
2. 아키텍처 범위: 본 연구 결과는 파라미터 공유 트랜스포머에 국한됩니다. 독립적인 레이어 가중치를 가진 표준 트랜스포머는 본질적으로 이러한 형태의 암묵적 일반화를 어려워하며, 결과가 해당 아키텍처로 확장되지 않을 수 있습니다.