LayerTrim이 어텐션은 지키고 MLP만 잘라 추론 모델을 압축한다
Unlocking Structured Pruning for Reasoning Language Models via Attention-Preserving MLP Trimming
무엇인가
추론 언어 모델(RLM)은 체인오브소트로 긴 사고 과정을 생성하면서 복잡한 문제를 푸는데, 그만큼 추론 비용이 커진다. 구조적 프루닝은 모듈을 통째로 제거해 더 작은 dense 모델을 만드는 방식으로 비용을 줄이지만, 기존 기법을 RLM에 그대로 적용하면 AIME 같은 장문 추론 과제에서 성능이 붕괴한다. 논문은 이 실패를 그림 1로 보여준다. 기존 프루닝 모델은 추론 흔적이 오히려 길어지면서 정확도는 떨어진다.
어떻게 동작하나
저자들은 Qwen3-14B에서 프루닝 전략을 분해해 세 가지 관찰을 얻는다. 첫째, MLP 트리밍이 어텐션 프루닝이나 모듈 전체 제거보다 표현 오차가 작다. 둘째, MLP 특징 중요도는 일부 특징에 집중된다. 셋째, 같은 예산이라도 모든 레이어에 고르게 분산하는 것보다 중복 레이어에 몰아서 트리밍할 때 오차가 작다. 이 관찰이 어텐션과 디코더 깊이는 그대로 두고 MLP 폭만 선택적으로 줄이는 설계의 근거다.
무엇과 다른가
LayerTrim의 절차는 다섯 단계다. 먼저 dense 모델이 생성한 추론 흔적에 대해 teacher-forced 순전파를 돌려 레이어별 입출력과 MLP 중간 활성값 통계를 모은다. 다음으로 Block Influence(BI = 1 − E[cos(h_in, h_out)])가 낮은, 즉 블록을 지나도 방향 변화가 작은 레이어를 고른다. 예산 N 레이어 상당에 대해 K = 2N개의 최저 BI 레이어를 선택해 그 MLP만 건드린다. 각 MLP에서는 SwiGLU 중간 특징 j가 출력에 기여하는 항 c_j(x) = φ_j(x)·W_down[:,j]의 기대 제곱 노름을 중요도 s_j로 정의하고, 상위 m개 특징만 남긴다. m은 3CK(I−m) = N·P_layer와 K = 2N에서 m = round(I − P_layer/6C)로 정해진다. 마지막으로 gate·up 투영의 행과 down 투영의 열을 물리적으로 잘라내 압축 체크포인트를 만든다. 어텐션 모듈, 선택되지 않은 MLP, 디코더 깊이는 손대지 않으며 그래디언트, 파인튜닝, 증류, 복구 학습, 추론 시 마스크가 모두 필요 없다.
어떻게 쓰나
실험은 Qwen3-14B, Phi-4-reasoning-plus, Magistral-Small-2506 세 모델과 AIME2025, AIME2026, GPQA-Diamond, HumanEval, MBPP+ 다섯 벤치마크에서 이뤄졌다. 깊이 프루닝 베이스라인은 ShortGPT, LaCo, GRASP, 폭 프루닝 베이스라인은 FLAP, RESP, CoCurve이며 예산은 7개와 9개 레이어 상당으로 파라미터 수를 맞춰 비교했다. 두 예산 모두에서 LayerTrim이 세 모델 모두 최고 평균 점수를 냈고, 예산이 커질수록 격차도 벌어졌다. N=9에서 가장 강한 베이스라인 대비 평균 우위는 Qwen3-14B 6.8%p, Phi-4-reasoning-plus 14.3%p, Magistral-Small-2506 21.4%p다. 캘리브레이션은 AIME 2024 추론 흔적만 썼는데도 GPQA-Diamond에서 최대 20.2%p, HumanEval 48.4%p, MBPP+ 35.0%p까지 이득이 번졌다.
전제와 한계
효율 수치도 함께 제시된다. 높은 프루닝 예산에서 프리필 속도는 세 모델 각각 1.24배, 1.22배, 1.27배, 디코드 속도는 1.21배, 1.22배, 1.24배가 됐고 최대 GPU 메모리는 19.7%, 20.3%, 20.8% 줄었다. Qwen3-14B 기준 프루닝 준비 시간은 27.9초로 RESP보다 N=7에서 94배, N=9에서 140배 빠르다. 분석 파트에서는 MLP 트리밍이 어텐션 KL 발산을 가장 낮게 유지하고, dense 모델의 top-8 어텐션 토큰과의 겹침도 16K~32K 거리 구간까지 가장 높게 보존한다. 특징 중요도 합으로 출력 변화를 근사하는 대리 지표는 Spearman ρ = 0.999로 측정값 순위를 재현하면서 캘리브레이션을 13.2배 빠르게 했다. 출력 진단에서는 Phi-4-reasoning-plus의 GPQA-Diamond 응답에서 토큰 한도 도달률 24.4%, 강한 반복률 21.9%로 FLAP의 70.2%, 49.2%보다 훨씬 낮았고, LaCo는 토큰 한도 도달 100%에 반복은 0%로 잡히지 않는 비정상 출력을 보였다. HumanEval에서 컴파일 가능한 코드 비율은 99.6%로 dense의 99.4%에 근접했고 RESP는 85.0%였다. AIME26 산술 일관성은 94.8%로 dense 96.4%에 1.6%p 차이까지 따라붙었다.
개발자 입장에서 이 논문의 실용적 요점은 프루닝 예산을 어디에 쓰느냐가 정확도를 좌우한다는 것이다. 어텐션과 레이어 깊이를 건드리지 않고 저 BI 레이어의 MLP 폭만 줄이면, 재학습 파이프라인 없이도 추론 모델을 압축할 수 있고 생성 토큰 수까지 줄어든다. 다만 캘리브레이션에 자체 추론 흔적이 필요하고, BI와 특징 중요도 통계가 모델·태스크에 따라 달라질 수 있으므로 배포 전에 자신의 워크로드에서 정확도 유지와 토큰 한도 도달률을 함께 확인해야 한다. K = 2N 기본값도 절대적이지 않다. 어블레이션에서 K = 1.5N은 AIME25, K = 3N은 GPQA-Diamond에서 각각 최고였고 K = 2N이 평균에서 근소하게 앞섰다.
저자들이 밝힌 한계는 분명하다. 평가가 세 개 추론 모델과 다섯 벤치마크(수학, 과학, 코드 생성)에 한정돼 있어 다른 아키텍처, 규모, 추론 과제로의 일반화는 추가 실험이 필요하다. 또한 프루닝 예산과 레이어 할당을 고정된 설정으로만 평가했으며, 모델과 태스크별로 이 선택을 적응시키면 압축과 정확도의 트레이드오프가 더 나아질 수 있다고 본다. 윤리 고려 사항에서는 프루닝된 모델이 여전히 틀린 답이나 불완전한 추론 흔적, 예상 밖 출력을 낼 수 있으므로 정확도가 중요한 응용에서는 위험을 감안해야 한다고 적었다.