기술적 요약: 인과적 에너지 최소화를 통한 트랜스포머 레이어 파라미터화 재검토

문제 제기

트랜스포머 아키텍처, 특히 멀티헤드 어텐션 (MHA) 과 게이트드 멀티레이어 퍼셉트론 (MLP) 의 조합은 현대 대규모 언어 모델 (LLM) 의 기본 구성 요소로 작용합니다. 그러나 레이어 파라미터화에 관한 구체적인 선택들 (가중치 공유, 투영 구조, 업데이트 규칙 등) 은 대부분 명시적인 계산 원리보다는 경험적 연구에 의해 주도되어 왔습니다. 에너지 기반 모델 (EBM) 은 과거에 계산을 저에너지 상태 탐색으로 바라보는 관점을 제시해 왔습니다 (예: 홉필드 네트워크). 하지만 기존 연구는 이러한 에너지 기반 해석을 표준 트랜스포머 레이어의 구체적인 파라미터화와 완전히 연결하지 못했으며, 이 관점이 더 효율적이거나 효과적인 레이어 변형의 설계에 어떻게 기여할 수 있는지를 체계적으로 탐구하지는 못했습니다.

방법론: 인과적 에너지 최소화 (CEM)

저자들은 인과적 에너지 최소화 (Causal Energy Minimization, CEM) 라는 프레임워크를 도입하여 트랜스포머 레이어를 조건부 에너지 함수 상의 최적화 단계로 재해석합니다. 핵심 전제는 레이어가 입력 시퀀스 $h_{1:J}$ 를 출력 $h'_{1:J}$ 로 매핑할 때, 각 위치 $i$ 를 초기값이 $h_i$ 인 최적화 변수 $x_i$ 로 간주하고, 인과적 역사에 의존하는 조건부 에너지 $\epsilon(x_i | h_{1:i})$ 를 최소화하기 위해 절차 $\mathcal{A}$ 를 통해 이를 업데이트한다는 것입니다.

이 논문은 이 프레임워크에서 표준 트랜스포머 구성 요소를 유도하고 확장을 제안합니다:

1. 가중치 공유 어텐션의 유도:
   현대 홉필드 네트워크와 유사한 로그 - 합 - 지수 (log-sum-exp) 형태를 기반으로 한 상호작용 에너지를 정의함으로써, 저자들은 단일 경사 하강 단계를 수행하면 표준 MHA 업데이트가 도출됨을 보여줍니다. 중요한 점은 이 유도가 특정 가중치 공유 제약을 자연스럽게 요구한다는 것입니다:

   • 키 (Key) 투영 ($W^K$) 은 값 (Value) 투영 ($W^V$) 과 공유됩니다.
   • 쿼리 (Query) 투영 ($W^Q$) 은 출력 (Output) 투영 ($W^O$) 과 공유됩니다.
   • 온도 매개변수 $\tau$ 는 헤드 차원과 관련됩니다.

2. 가중치 공유 게이트드 MLP 의 유도:
   각 토큰의 특징 벡터에 독립적으로 에너지를 할당하는 원소별 에너지를 정의하면, 경사 업데이트가 게이트드 MLP 의 구조를 복원합니다. 이 해석은 MLP 내의 "업 (up)" 및 "다운 (down)" 투영 행렬이 공유되어야 함을 시사합니다 (구체적으로, 다운 투영의 전치 행렬이 업 투영과 같아야 함).

3. 제안된 확장:
   에너지 최적화 관점을 바탕으로 저자들은 세 가지 구체적인 아키텍처 향상을 제안합니다:

   • 대각선 + 저랭크 (D+LR) 파라미터화: 상호작용 행렬 $A_k$ 에 대해 순수한 저랭크 투영 대신, $A_k = \text{diag}(d_k) + W^{Q\top}_k W^K_k$ 를 제안합니다. 이는 저랭크 행렬이 놓치는 상호작용을 포착하면서도 계산 효율성을 유지합니다. 또한 파라미터를 줄이기 위해 헤드 간 대각선 항을 공유하는 것을 제안합니다.
   • 학습된 경량 사전 조건부 (Learned Lightweight Preconditioners): 단일 경사 단계 이상의 수렴을 개선하기 위해, D+LR 구조 ($P = \text{diag}(\text{softplus}(d)) + UV^\top + VU^\top$) 를 가진 사전 조건부 행렬 $P$ 를 도입합니다. 이들은 곡률 정보의 학습된 대리자로 작용하여 완전한 2 차 방법의 비용 없이 경사를 재조정합니다.
   • 레이어 내 재귀적 업데이트: 단일 업데이트 대신, 프레임워크는 단일 레이어 내에서 $T$ 단계의 재귀적 업데이트를 허용합니다 ($x^{(t+1)}_i = x^{(t)}_i - \eta \nabla \epsilon$). 이는 새로운 파라미터를 추가하지 않고도 레이어가 에너지 함수를 더 깊게 최소화할 수 있게 하며, 전체 블록을 재사용하는 "루프형" 트랜스포머와는 구별됩니다.

주요 기여

• 이론적 통합: 이 논문은 특정 가중치 공유 방식 (어텐션에서 Q/O 및 K/V 공유; MLP 에서 업/다운 공유) 을 가진 표준 트랜스포머 레이어가 특정 상호작용 및 원소별 에너지 함수 상의 단일 경사 단계와 수학적으로 동등함을 확립합니다.
• 설계 공간 확장: CEM 관점은 D+LR 상호작용, 경량 사전 조건부, 재귀적 업데이트를 포함한 트랜스포머 레이어의 새로운 설계 공간을 식별하며, 이는 휴리스틱이 아닌 최적화 이론에 기반합니다.
• 경험적 검증: 저자들은 언어 모델링 작업에서 수억 개의 파라미터 규모로 CEM 에서 유도된 레이어를 평가하여, 이러한 제약된 파라미터화가 표준 트랜스포머 베이스라인과 경쟁하거나 근접할 수 있음을 입증했습니다.

실험 결과

실험은 SlimPajama 데이터셋에서 약 86M 에서 약 162M 개의 파라미터 크기를 가진 모델로 수행되었으며, 컴퓨팅 최적화 토큰 예산 (Chinchilla 스케일링) 으로 학습되었습니다.

• 파라미터 효율성: 가중치를 공유한 CEM 모델 (Llama 베이스라인 대비 MHA 파라미터의 약 절반, MLP 파라미터의 약 3 분의 2 사용) 은 경쟁력 있는 테스트 퍼플렉시티를 달성했습니다. CEM 어텐션의 경우, 파라미터가 크게 감소했음에도 성능 저하는 미미했습니다.
• 확장의 영향:
  • 재귀: 재귀 단계 수를 $T=1$ 에서 $T=2$ 로 늘리는 것이 CEM 어텐션과 MLP 모두에서 일관되게 성능을 향상시켰습니다. 특히, 재귀적 CEM 어텐션 변형은 더 적은 파라미터를 사용하면서도 Llama 베이스라인을 능가했습니다.
  • 사전 조건부: 학습된 사전 조건부는 재귀에 비해 미미한 개선만 제공하여, 재귀 구조가 개선의 주된 동인임을 시사합니다.
  • 대각선 항: 상호작용 행렬에 대각선 항을 포함하는 것 (특히 헤드 간 공유되는 대각선) 은 양호한 성능에 필수적이었으며, 순수한 저랭크 또는 대각선 항이 없는 변형보다 우수한 성능을 보였습니다.
• 안정성: 완전히 CEM 에서 유도된 레이어로 구성된 트랜스포머 ($T=2$ 및 D+LR 사전 조건부 사용) 는 엔드 - 투 - 엔드 학습에서 안정적으로 학습되어, 더 적은 파라미터로 Llama 베이스라인과 비교 가능한 퍼플렉시티를 달성했습니다.
• 합성 데이터: 합성 가우스 프로세스 데이터에 대한 실험은 더 깊은 재귀 ($T=4, T=8$) 가 RMSE 에서 추가적인 개선을 가져온 것을 확인시켰으나, 표준 학습 설정 하의 전체 트랜스포머 환경에서는 $T=2$ 까지의 개선으로 제한되었습니다.

중요성과 주장

이 논문은 CEM 이 트랜스포머 레이어 파라미터화를 이해하는 데 유용한 이론적 렌즈를 제공하며, 트랜스포머 아키텍처를 에너지 기반 모델과 효과적으로 연결한다고 주장합니다. 그 중요성은 다음과 같습니다:

1. 원칙적 설계: 레이어 설계를 경험적 시행착오에서 에너지 최소화와 최적화 역학에 기반한 프레임워크로 전환합니다.
2. 효율성: 이 프레임워크에서 유도된 가중치 공유와 재귀적 업데이트가 경쟁력 있는 성능을 유지하면서 파라미터 효율적인 모델을 가능하게 함을 입증합니다.
3. 향후 방향: 저자들은 이 관점이 다음과 같은 추가 탐구를 동기부여한다고 제안합니다:
   • 커스텀 커널: 구조화된 가중치 공유 및 재귀적 업데이트에 대한 런타임 최적화 (예: 연산 융합).
   • 테스트 시간 스케일링: 추론 시 컴퓨팅을 확장하기 위한 추가 축으로서 레이어 내 재귀 활용.
   • 하드웨어 공동 설계: 새로운 하드웨어 아키텍처에 맞춰 레이어 파라미터화 설계.

저자들은 평가가 수억 개의 파라미터 규모로 제한되었으며, 이러한 이점의 최첨단 모델로의 전이 가능성을 완전히 평가하기 위해서는 더 큰 규모의 연구와 하류 작업 평가가 필요하다고 겸손하게 언급합니다. 또한 현재 구현체는 런타임 성능에 최적화되지 않았음을 인정합니다.