모두 설명하는 한 층: 대형 언어 모델에서 Massive Activations의 기원과 전개

대형 언어 모델에서 일부 토큰의 활성화가 다른 토큰에 비해 수십에서 수백 배 크게 증가하는 현상(Massive Activations)이 관찰된다. 이 현상은 모델의 표현 다변성 감소와 self-attention의 입력 의존성 저하로 이어질 수 있으며, ME Layer의 구체적 기전과 그 후속 효과를 이해하는 것이 모델 해석 및 개선에 핵심이다. 본 논문은 ME Layer를 공통의 계층으로 확인하고 RMSNorm와 FFN이 결합하여 대량 활성화를 유발한다는 근거를 제시하며, 이로 인한 방향성 고정성을 완화하는 WeMask를 제안한다. 이 접근은 다양한 모델 군에서 일관된 성능 향상을 보이고, attention sinks의 기원에 대한 새로운 관점을 제공한다.

단계 1: 특정 ME Layer에서 토큰 단위의 대량 활성화가 급격히 증가하고, 이후 잔류 연결로 안정적으로 전파된다. 단계 2: ME Layer의 RMSNorm 가중치와 FFN의 프로젝션이 이 증폭을 주도하며, 첫 토큰의 활성화가 다른 토큰에 비해 크게 증가한다. 단계 3: ME Layer 이후에는 이 대량 활성화가 방향적으로 안정되어 다른 입력에서도 비슷한 방향성을 보이고, self-attention의 키/질의가 거의 고정된 방향으로 작용하게 된다. 단계 4: WeMask는 pre-attention RMSNorm의 큰 가중치를 가진 차원을 마스크해 방향성의 고정을 완화하고, 그러면서도 전체 구조를 유지하여 성능을 향상시킨다.

전체 흐름은 ME Layer의 위치 규명 → RMSNorm와 FFN의 기여 분석 → ME Layer 이후의 방향성 분석 → WeMask 설계 및 구현 순으로 전개된다. 수식/수학적 흐름은 아래와 같다: RMSNorm의 출력은 x의 각 차원에 대해 x_i^2의 평균에 epsilon를 더해 제곱근을 취하고, learnable scale w를 곱해 차원별 재가중을 수행한다. 토큰 t의 RMSNorm 출력 h_hat_t,i는 h_t,i의 제곱 합에 의거한 가중치로 재배치되며, 첫 토큰의 분포가 큰 가중치 차원에 더 집중된다. FFN은 입력 벡터 h를 특정 차원에 집중시키는 투영(UpProj, GateProj, DownProj)을 통해, ME Layer에서 강한 magnitude 증폭을 유도한다. 이때 ME Layer의 증폭은 레이어7에서 가장 크게 나타나고, 이후 레이어로 전달될수록 magnitude는 유지되지만 방향은 거의 변하지 않는다. WeMask는 S(l) = TopK(w(l))으로 RMSNorm 가중치가 큰 차원을 선택하고, m(l)_d를 1 또는 0으로 설정하여 h̃^(l)_0 = h^(l)_0 ⊙ (1 − m(l))로 attention 입력 차원을 마스킹한다. 이 모듈은 ME Layer 이후의 모든 layer에 적용 가능하며, training-free 및 fine-tuning 설정에서 성능을 개선한다.

주요 벤치마크 결과는 다음과 같다. Instruction fine-tuning에서 Qwen3-4B + SFT에 WeMask(TF) 적용 시 MMLU는 54.32→54.32(0.1 비율) 등 소폭 개선, PIQA 80.69→81.03, ARC-C 87.54→86.25~87.54 범위, MathQA 37.76 등으로 나타났고 합산 평균도 증가했다. 또한 WeMask(SFT) 적용 시 0.7 비율에서 AVG가 64.88로 최적화되었다. 수치 중 일부는 마스크 비율에 따라 달라지며, 1.0으로 지나치게 높일 경우 성능이 감소하는 경향이 관찰된다. Math reasoning 및 safety alignment 태스크에서도 WeMask의 training-free/finetune 설정이 성능 향상을 이끌었으며, XSTest의 거부 경향을 완화하면서도 전반적인 성능을 유지하거나 개선하는 경향이 확인된다. RL 기반 학습(DPO, GRPO)에서도 WeMask를 도입하면 XSTest에서 성능이 상승하고 AdvBench에서의 비용 지표는 감소하는 경향이 나타나며, GRPO + WeMask(TF)에서 수학 문제(Math500) 및 수력(AIME/22–24)에서도 향상을 보인다. 모델 계열 간(Measure across Llama3-8B-Instruct, Qwen-8B, 등) ME Layer 위치는 대체로 유사하며, Qwen3-4B-Instruct/8B의 ME Layer 위치는 7번으로 일관된다. 또한 WeMask는 다양한 모델 유형에서 일반화되며, SFT 기반 모델에서도 일관된 개선을 보인다.

1. 아키텍처 구성: Transformer Decoder 블록에서 ME Layer 이후의 각 layer에 WeMask 모듈을 삽입한다. 2) 핵심 메커니즘: RMSNorm 가중치가 큰 차원을 TopK로 선택하고, 해당 차원을 입력 숨겨진 상태에서 마스크한다. 마스크된 차원은 attention 입력에서 제거되어, 키-질의 방향성이 더 다양해지도록 한다. 3) 차별점: FFN이 대량 활성화를 생성하는 주된 원인으로 확인되지만, pre-FFN RMSNorm의 가중 재조정이 이 활성화의 규모를 좌우한다. 4) 구현 및 학습 세부사항: 학습 중 ME Layer 이후 모든 layer에 WeMask를 적용해 representation-level regularization 효과를 유도하고, inference 시 태스크에 따라 ME Layer 이후의 일부 layer에만 적용하는 전략도 제시한다.