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

Revisiting Transformer Layer Parameterization Through Causal Energy Minimization

트랜스포머 블록은 일반적으로 토큰 믹싱을 위한 멀티 헤드 어텐션(MHA)과 토큰 단위 특징 변환을 위한 게이티드 MLP를 결합하지만, 이러한 파라미터 설정 방식은 대부분 경험적인 수준입니다. 본 연구에서는 트랜스포머 레이어를 조건부 에너지 함수에 대한 최적화 단계로 재해석하고 레이어 파라미터화를 명시적으로 고려하는 프레임워크인 인과 에너지 최소화(CEM)를 제안합니다. 기존의 어텐션에 대한 에너지 기반 해석을 확장하여, CEM은 가중치 공유 MHA가 상호 작용 에너지에 대한 기울기 업데이트로 유도될 수 있으며, 공유 업/다운 프로젝션을 가진 게이티드 MLP는 요소별 에너지 관점에서 볼 수 있음을 보여줍니다. 이러한 관점은 레이어 내 가중치 공유, 대각선-저랭크 상호 작용, 경량 프리컨디셔너 및 재귀적 업데이트를 포함하는 트랜스포머 레이어 설계 공간을 제시합니다. 우리는 중간 규모(수억 파라미터)의 언어 모델링 실험에서 CEM에서 파생된 레이어를 평가했습니다. 이러한 제한된 파라미터화에도 불구하고, 이러한 레이어는 안정적으로 학습되며 기존 트랜스포머 모델과 유사한 성능을 보입니다. 전반적으로, 본 연구의 결과는 CEM이 트랜스포머 레이어 파라미터화에 대한 유용한 이해를 제공하며, 트랜스포머 아키텍처를 에너지 기반 모델과 연결하고 에너지 기반 레이어 설계에 대한 추가적인 연구를 촉진할 수 있음을 시사합니다.

Transformer blocks typically combine multi-head attention (MHA) for token mixing with gated MLPs for token-wise feature transformation, yet many choices in their parameterization remain largely empirical. We introduce Causal Energy Minimization (CEM), a framework that recasts Transformer layers as optimization steps on conditional energy functions while explicitly accounting for layer parameterization. Extending prior energy-based interpretations of attention, CEM shows that weight-tied MHA can be derived as a gradient update on an interaction energy, and that a gated MLP with shared up/down projections can be viewed through an element-wise energy. This perspective identifies a design space for Transformer layers that includes within-layer weight sharing, diagonal-plus-low-rank interactions, lightweight preconditioners, and recursive updates. We evaluate CEM-derived layers in language-modeling experiments at the moderate hundred-million-parameter scale. Despite their constrained parameterizations, these layers train stably and can match corresponding Transformer baselines. Overall, our results suggest that CEM provides a useful lens for understanding Transformer layer parameterization, connecting Transformer architectures to energy-based models and motivating further exploration of energy-guided layer designs.