Half the Nonlinearity Is Wasted: Measuring and Reallocating the Transformer's MLP Budget

이 논문은 트랜스포머의 MLP 비선형성이 토큰 정체성이 아닌 문맥에 따라 결정되며, 게이트 메커니즘을 통해 대부분의 계산을 선형으로 대체하거나 특정 계층을 제거함으로써 퍼플렉시티 손실 없이 모델 성능을 오히려 향상시킬 수 있음을 입증합니다.

핵심

🍔 햄버거 가게의 비유: "과도한 요리사"

생각해 보세요. 거대한 햄버거 가게 (AI 모델) 가 있습니다. 이 가게에는 매번 주문이 들어올 때마다 **매우 정교하고 복잡한 요리를 해내는 요리사 (MLP, 비선형 회로)**가 있습니다.

• 기존의 생각: "모든 햄버거 (단어) 를 만들 때, 이 요리사의 복잡한 기술이 꼭 필요하다. 만약 요리사를 없애거나 단순하게 만들면 햄버거 맛이 망가질 것이다."
• 이 논문의 발견: "아니요! 사실 **대부분의 주문은 그냥 빵과 패티를 쌓는 것 (선형 계산)**만으로도 충분합니다. 요리사가 복잡한 소스를 뿌리거나 고기를 다지는 과정은 90% 이상의 경우 불필요한 낭비입니다."

연구진은 이 "낭비되는 요리 시간"을 찾아내고, 어떤 경우에는 복잡한 요리사를 쉴게 하고, 간단한 조리대 (선형 행렬) 만으로 대신하게 하는 시스템을 만들었습니다.

---

🔍 주요 발견 3 가지

1. "누가 요리사가 필요할까?"를 단어만으로 판단할 수 없다

처음 연구진은 "아마도 '사과', '자동차' 같은 구체적인 단어 (명사) 는 복잡한 요리가 필요하고, '의', '은' 같은 조사 (기능어) 는 간단하게 처리하면 되겠지?"라고 생각했습니다.

하지만 결과는 완전한 실패였습니다.

• 비유: "사과"라는 단어가 들어와도, 문맥에 따라서는 그냥 껍질만 벗겨주면 될 때도 있고 (간단), 껍질을 벗기고 씨를 빼고 조각내야 할 때도 (복잡) 있습니다.
• 결론: 단어 자체의 종류로는 누가 복잡한 계산이 필요한지 알 수 없습니다. 오직 그 단어가 어떤 문맥 (상황) 에 놓여 있는지를 봐야만 알 수 있습니다.

2. "스마트한 문지기 (게이트)"를 도입하자

연구진은 아주 간단한 **문지기 (게이트)**를 두었습니다. 이 문지기는 "이 주문이 복잡한 요리가 필요한가, 아니면 간단한 조리로 충분할까?"를 0.1 초 만에 판단합니다.

• 결과: 이 문지기가 판단해서 약 40~50% 의 계산은 복잡한 요리사 없이 간단한 조리대로 처리했습니다.
• 효과: 놀랍게도 햄버거의 맛 (모델의 성능) 은 거의 떨어지지 않았습니다. 오히려 일부 층에서는 복잡한 요리사가 오히려 맛을 망쳤기 때문에, 간단한 조리대로 바꾸니 맛이 더 좋아지기도 했습니다.

3. 모델의 '건축 방식'이 중요했다

이 실험을 두 가지 다른 모델 (GPT-2 와 Pythia) 에 해보았는데 결과가 달랐습니다.

• GPT-2: 대부분의 층에서 복잡한 요리사를 줄여도 문제가 없었습니다.
• Pythia: 처음과 마지막 층은 요리사가 꼭 필요했지만, 중간 층은 GPT-2 보다는 조금 더 복잡했습니다.
• 교훈: 모델의 내부 구조 (건축 방식) 에 따라 낭비되는 계산의 양이 다릅니다.

---

🚀 이 연구가 우리에게 주는 메시지

1. 낭비는 실존한다: 우리가 AI 에게 "모든 상황에서 최고의 성능을 내라"고 강요하다 보니, 실제로는 필요 없는 복잡한 계산이 대량으로 수행되고 있었습니다.
2. 상황이 모든 것이다: "이 단어가 중요해"라고 미리 정해둘 수 없습니다. 오직 **지금의 상황 (문맥)**을 보고 판단해야 합니다.
3. 미래의 AI 설계: 앞으로의 AI 는 모든 층에 똑같은 고성능 엔진을 달지 않아도 됩니다.
   • 입구와 출구 (처음과 마지막): 고성능 엔진 (복잡한 비선형 계산) 을 달아줍니다.
   • 중간 층: 효율적인 소형 엔진 (간단한 선형 계산) 으로 교체합니다.
   • 결과: 더 적은 전력과 비용으로 똑같거나 더 좋은 성능을 낼 수 있습니다.

💡 한 줄 요약

"AI 가 모든 단어를 처리할 때 '최고급 요리사'를 쓸 필요는 없습니다. 대부분의 경우 '간단한 조리대'로 충분하며, 상황에 따라 문지기가 이를 잘 판단하게 하면 AI 는 더 빠르고 똑똑해집니다."

기술 요약

논문 요약: "Half the Nonlinearity Is Wasted: Measuring and Reallocating the Transformer's MLP Budget"

이 논문은 트랜스포머 (Transformer) 모델의 핵심 구성 요소인 다층 퍼셉트론 (MLP) 의 비선형성 (nonlinearity) 이 실제로 얼마나 필요한지, 그리고 그 계산 자원을 어떻게 최적화할 수 있는지에 대한 체계적인 연구를 제시합니다. 저자는 "비선형성의 절반은 낭비된다"는 가설을 검증하여, 많은 MLP 계산을 선형 행렬로 대체하거나 게이트 (gate) 를 통해 동적으로 라우팅함으로써 성능 저하 없이 계산 비용을 크게 절감할 수 있음을 증명했습니다.

---

1. 연구 배경 및 문제 정의 (Problem)

• 기존 가정: 트랜스포머 아키텍처에서 각 레이어의 MLP 는 복잡한 비선형 변환을 수행하며, 이 비선형성 (활성화 함수) 은 언어 모델링에 필수불가결하다고 간주되어 왔습니다. 비선형성을 제거하면 네트워크가 단순한 선형 매핑으로 축소되어 복잡한 언어 기능을 학습하지 못한다고 믿어집니다.
• 문제 제기: 모든 토큰과 모든 레이어에서 정교한 비선형 계산을 수행하는 것이 정말로 필요한가? 만약 많은 계산이 사실상 선형에 가깝다면, 이 자원을 어떻게 재분배하거나 줄일 수 있는가?
• 목표: MLP 비선형성의 실제 필요성을 측정하고, 이를 대체하거나 라우팅하는 효율적인 방법을 모색하여 모델의 계산 효율성을 높이는 것.

2. 방법론 (Methodology)

저자는 6 개의 모델 (GPT-2 Medium/Large, Pythia 시리즈 등, 1.62 억~28 억 파라미터), 2 가지 아키텍처, 3 개의 코퍼스를 대상으로 다음과 같은 실험을 수행했습니다.

1. 선형 근사 (Linear Approximation):

   • 각 레이어의 MLP 입력과 출력을 수집하여 리지 회귀 (Ridge Regression) 를 통해 최적의 선형 행렬 ($Wx + b$) 을 구했습니다.
   • 전체 MLP 를 이 선형 대용품으로 교체했을 때의 퍼플렉시티 (Perplexity, PPL) 증가분을 측정하여 '선형화 비용'을 평가했습니다.

2. 적응형 게이팅 (Adaptive Gating):

   • 모든 토큰에 대해 MLP 를 사용하는 대신, 게이트 (Gate) 가 각 토큰의 활성화 값을 보고 '선형 대용품'을 사용할지 '전체 MLP'를 사용할지 결정하는 방식을 도입했습니다.
   • 게이트 구조: 가장 간단한 로지스틱 회귀 (d+1 파라미터, 여기서 d 는 히든 차원) 를 사용했습니다.
   • 학습 데이터: 토큰별 손실 차이 ($\delta = L_{lin} - L_{full}$) 를 기반으로 '선형으로 대체 가능한지'를 이진 분류 문제로 학습했습니다.

3. 라우팅 신호 분석 (Routing Analysis):

   • 토큰 정체성 vs. 문맥: 게이트가 어떤 신호를 기반으로 결정을 내리는지 분석하기 위해, MLP 입력을 '토큰 임베딩 (정체성)'과 '문맥 기여도 (Attention 및 이전 레이어의 결과)'로 분해하여 각각 게이트를 학습시켰습니다.
   • 교차 코퍼스 테스트: 한 코퍼스 (WikiText-103) 에서 학습한 '비선형성이 필요한 토큰 목록 (No-Fly lists)'을 다른 코퍼스 (LAMBADA 등) 에서 테스트하여 일반화 능력을 검증했습니다.

4. 점진적 선형화 (Progressive Linearization):

   • 여러 레이어의 MLP 를 동시에 선형 행렬로 교체하고, 나머지 레이어를 미세 조정 (Fine-tuning) 하여 모델이 적응할 수 있는지 확인했습니다.

3. 주요 기여 및 결과 (Key Contributions & Results)

3.1. MLP 계산의 대부분은 사실상 선형임

• 결과: 대부분의 모델 (특히 GPT-2) 에서 중간 레이어의 MLP 를 완전히 선형으로 대체해도 퍼플렉시티 손실이 미미했습니다 (대부분 3% 미만).
• 특이점: GPT-2 Medium 의 경우 23 개 레이어 중 4 개 레이어는 오히려 비선형성을 제거했을 때 성능이 향상되었습니다. 이는 해당 레이어의 MLP 가 과적합 (Overfitting) 을 유발하고 있었음을 의미합니다.

3.2. 토큰 기반 라우팅의 실패 (Strong Negative Result)

• 핵심 발견: "어떤 토큰이 비선형성이 필요한가?"를 토큰의 정체성 (예: 기능어 vs 내용어) 으로 예측하는 것은 불가능합니다.
• 데이터: 한 코퍼스에서 학습한 '비선형성 필요 토큰 목록'을 다른 코퍼스나 같은 도메인의 다른 텍스트에 적용했을 때 상관관계가 거의 0 에 가까웠습니다 ($r < 0.05$).
• 결론: 비선형성 필요 여부는 토큰 자체의 속성이 아니라, **문맥 (Context)**에 의해 결정됩니다. 같은 단어라도 문맥에 따라 선형 처리가 필요할 수도 있고 비선형 처리가 필요할 수도 있습니다.

3.3. 문맥 기반 게이팅의 성공

• 게이트 성능: 토큰 정체성만 보는 게이트는 무작위 수준이었으나, **문맥 기여도 (Contextual contribution)**만 보는 게이트는 전체 게이트와 거의 동일한 성능을 냈습니다.
• 효율성: 단순한 로지스틱 회귀 (d+1 파라미터) 게이트만으로도 GPT-2 Medium 에서 25~56% 의 MLP 계산을 선형 경로로 라우팅하면서도 퍼플렉시티 손실은 1% 미만으로 유지했습니다.
• 성능 향상: GPT-2 Medium 의 4 개 레이어에서는 게이트를 통해 선형 경로를 선택했을 때 오히려 베이스라인보다 퍼플렉시티가 개선되었습니다.

3.4. 아키텍처 의존성 (Architecture Dependence)

• GPT-2 vs Pythia: GPT-2 아키텍처는 선형화가 매우 용이한 반면, Pythia (GPT-NeoX 기반) 는 선형화 비용이 더 높았습니다.
• 규모의 영향: Pythia-2.8B 의 경우 중간 레이어 (L7-L15) 는 선형화 비용이 낮았으며, L3 레이어는 베이스라인을 약간 상회하는 성능을 보였습니다. 하지만 Pythia-2.8B 의 첫 번째 레이어 (L0) 를 선형화하면 모델이 붕괴되는 (+513% PPL 증가) 치명적인 차이가 있었습니다. 이는 아키텍처 설계 (순차적 vs 병렬적 계산) 가 비선형성 필요도에 큰 영향을 미친다는 것을 보여줍니다.

3.5. 점진적 선형화 및 미세 조정 (Proof of Concept)

• 실험: GPT-2 Medium 의 중간 레이어 5 개를 선형 행렬로 교체하고 미세 조정했습니다.
  • 결과 1: 5 개 레이어를 선형화해도 퍼플렉시티 변화가 0% 였습니다.
  • 결과 2: 4 개 레이어를 선형화하고 충분한 데이터 (1.17 억 토큰) 로 미세 조정했을 때, 원래 모델보다 10.2% 퍼플렉시티가 개선되었습니다.
  • 결과 3: 2 단계 게이팅 접근법 (선형화 + 학습된 게이트) 을 적용하면 **17.3% 개선 (PPL 19.00)**을 달성하여, 동일한 계산 비용의 일반 미세 조정 제어군보다도 성능이 뛰어났습니다.

4. 의의 및 시사점 (Significance)

1. 비선형성 자원의 재평가: 트랜스포머의 비선형성은 모든 레이어와 모든 토큰에 균일하게 필요한 것이 아니라, 문맥에 의존적이며 분포가 치우쳐 있음 (대부분의 계산은 선형에 가까움) 을 증명했습니다.
2. 효율적인 추론의 가능성: 토큰 기반의 라우팅 (예: 특정 단어는 생략) 은 불가능하지만, 문맥을 기반으로 동적으로 라우팅하거나 아키텍처 자체를 변경함으로써 계산 효율성을 극대화할 수 있습니다.
3. 차세대 아키텍처 설계 방향:
   • 비균일 용량 할당: 경계 레이어 (입력/출력) 에는 큰 비선형 MLP 를 배치하고, 중간 레이어는 선형 또는 작은 MLP 로 대체하는 설계가 유효합니다.
   • 하이브리드 레이어: 선형 잔차 연결과 작은 비선형 MLP 를 결합한 하이브리드 구조가 유망합니다.
   • 학습 중 라우팅: 사전 학습 단계에서 게이트와 MLP 가 함께 학습되도록 하여, 비선형성이 필요한 경우에만 자원을 할당하는 아키텍처가 제안됩니다.
4. 메커니즘 해석의 교훈: 단순한 토큰 정체성과의 상관관계만으로는 모델의 동작을 해석할 수 없으며, 교차 코퍼스 검증이 필수적임을 강조했습니다.

결론

이 논문은 "비선형성의 절반은 낭비된다"는 통찰을 통해, 트랜스포머 모델의 MLP 자원을 더 효율적으로 재분배할 수 있음을 실증했습니다. 특히 문맥 기반의 동적 라우팅과 아키텍처 수준의 비선형성 최적화를 통해, 파라미터 수를 줄이거나 계산 비용을 절감하면서도 오히려 성능을 향상시킬 수 있는 새로운 가능성을 제시했습니다.