A Simple Linear Patch Revives Layer-Pruned Large Language Models

NeurIPS 2025

👀 요약 👀


✨ method 정리 ✨
프루닝된 레이어 사이에 activation channel간 magnitude가 매우 불일치한 현상에 주목.

이 activation scale을 맞춰주기 위한 scaling factor를 도입한다.
 1. channel-wise scaling : d
      프루닝 이후 영향을 받는 두 레이어간의 activation (X)의 평균 activation magnitude의 비율
 2. token-wise scaling : H
      outlier가 되는 토큰들이 있다.(eg. [BOS] ...) 이를 완화하기 위해 Hadamard transform을 적용한다.

위 두 개의 scaling 과정을 하나로 결합하여 patch matrix P를 만든다. (dim x dim)
P는 offline distillation 과정으로 KL-div를 사용하여 finetuning 시킨다.

완성.

---
* 근데 pruning할 레이어 정하는 거는 그냥 기존에 많이 사용하는 대로 cosine sim을 썼다고 밝힘.
그럼 기존 방법에 새로운 레이어를 넣는 방법이니 당연하게도 좋아지지 않으려나 싶긴 함.
저 새로운 patch 만드는 방법이 현상(activation mag 불일치) 있어보이는 다른 방법(hardamard transform)을 끌어다와가지고 novelty가 생긴 것 같다.
* 그리고 튀는 엑티베이션이 있다면.. 그것조차 티처모델에서 나온건데 유지해야 하는 거 아닌가? 저 아다마르 변환은 그걸 유지하나?? 아니 레이어가 제거되지 않으면 그 튀는 엑티베이션이 점점 사라지는 거야? 그게 아니라면 굳이 없앨필요가 없잖아

---
목차부터 뭔가 깔끔하네
신박하긴 함
새로운 레이어.. 
어쨌든 finetuning을 시키는 거면 초기화에 불과한 것 같음. (좋은초기화? 얼마나 좋은데??)
그리고 뒷쪽 레이어의 아웃풋과 activation을 유사하게 맞춰서 넣어준다....라는 게 왜 작동을 하지??
원래 input으로 들어가는 activation이 비슷해서 그런건가???

 

Abstract로 흐름 파악하기

 

Layer pruning은 LLM을 compress하는데 widely하게 사용되는 방법이다.

하지만 기존 layer purning 방법들은 상당한 성능 저하가 발생한다.

 

본 논문에는 이러한 성능 저하의 대부분이 이전에는 간과되었던 프루닝 인터페이스에서의 activation magnitudes 불일치 문제에서 기인함을 확인했다.

프루닝 전후에 활성화되는 스케일(????)이 많이 달라져서, 남은 레이어를 거치면서 distributional shift가 일어난다.

different activation scale이 머임?????????????????? 설명 제대로 해주지 않으면 화가 날것.

 

이 문제를 해결하기 위해 LinearPatch를 제안,

lightweight 하고 plug-and-play한 방법이며, 프루닝 과정(interface)에서 두 개의 연산을 하나의 matrix multiply로 통합한다.

(i) 특정 토큰에서 발생하는 거대한 outlier들을 억제하기 위한 Hadamard transformation

오호 이부분 궁금함.. 거대한 outlier조차 원래 모델의 지식일텐데.. 데이터를 많이 뽑아서 사용하면 될 것 같은데? 의도적으로 제거하기?

거대한 outlier 데이터가 어느 calibration 또는 몇번째레이어에서 나오는지도 알려주나? 알려주겠지?

(ii) activation statistics를 정렬(align)하기 위한 channel-wise scaling

 

 

LaMA-3-8B 모델에서 LINEARPATCH는 32개 레이어 중 5개를 프루닝할 때도 94.15%를 유지하며, 이전 SOTA 방법 대비 4% 높은 성능. (5개는 ....  15% 프루닝한건데 ....... 나도 이렇게 자랑해야겠다...... )

5천 개의 라벨 없는 샘플을 활용한 메모리 효율적인 offine distillation으로 패치를 추가로 정제하면, 단일 GPU에서 30분 만에 성능 유지율을 95.16%까지 끌어올릴 수 있다.

---

1. Introduction

 

레이어프루닝이 emerge하고 있다. 특별한 하드웨어 specific한 optimization이나 low-level kernel modification에 의존하지 않는 방법이기 때문이다. 별도의 dependency 없이 불필요한 레이어를 제거하는 간단한 방식이다. 

반면 unstructured pruning은 불규칙적인 메모리 접근 패턴 때문에 가속화가 어렵고,

structured의 경우에는 종종 모델 아키텍처의 변형 또는 맞춤형 kernel이 요구된다는 문제가 있다. 

레이어프루닝은! 별도의 dependency 없이 불필요한 레이어를 제거하는 간단한 방식이다.  --> 하지만 성능 저하가 심하다는 크리티컬한 challenge들이 있다. 

이 연구에서는 이러한 성능 저하를 설명하는 새로운 현상을 발견하였다: 프루닝 지점에서 layer과 token 간의 activation magnitude 불일치이다.
구체적으로, 일부 층이 프루닝될 때 남은 층들의 activation 값은 종종 서로 다른 스케일을 보이며, 프루닝 지점 이전 층의 activation이 이후 층의 activation 정렬되지 않을 수 있다. 이러한 불일치는 특수 토큰(eg.: [BOS] 또는 구분자 토큰)의 활성화에서 관찰되는 극단적인 outlier 존재로 인해 더욱 심화된다. (ref1, ref2)
결과적으로 프루닝된 LLM은 심각한 activation 불일치를 겪게 되며, 이는 결국 성능 저하로 이어진다.

 

이러한 이슈를 제거하기 위해 LinearPatch 메서드를 제안한다. 위에서 언급한 activateion mismatch를 완화하기 위해 디자인된 plug-and-play 방법이다. LInearPatch는 다양한 pruning metric에 간단하게 적용될 수 있다.

우선 Hadamard transformation를 적용하여 스페셜토큰에 대한 activation (== outliers) 를 억제시킨다.

 

이후 channel-wise scaling parameter를 도입하여, activateion magnitude에 있는 갭을 메운다. Spectral Theory에 의하면 hardamard transformation과 diagonalized channel-wise scaling은 하나의 real symmetric matirx로 표현할 수 있으며, 이를 LinearPatch에 활용한다. (머라노다른지식뭐야이거)

이 방법은 추론(inference) 오버헤드를 거의 발생시키지 않으면서도 활성화 크기를 효과적으로 정렬(alignment)한다.

 

정렬 이외에도, memory-efficient knowledge distillation를 통해 가지치기된 LLM을 추가로 향상시킨다. 구체적으로, 모든 다른 모델 파라미터를 고정한 채 LINEARPATCH 행렬만 finetuning 한다. 단 5,000개 샘플만 사용해도 되며, 7B 규모 모델 기준으로 단일 GPU에서 30분 이내에 완료할 수 있다.

 

실험 결과~~

벤치마크에서 LLaMA-3-8B의 5개 층을 가지치기한 경우, LINEARPATCH는 기존 성능의 94.15%를 유지,

LLM-Streamline(90.84%) 등 최신 방법들을 크게 능가 (오호.)

 

 

 

2. Related Work

Weight Pruning

- (unstructured) Wanda

- (structured) entire groups of weights를 제거하는 방법 (attention heads, MLP neuraons, or hidden dimenstions)

   - N:M sparsity

   - unstructure purning보다는 하드웨어 friendly하지만, 역시 재학습이 필요하다는 문제가 있다.

 

 

Layer Pruning

레이어프루닝 등장

width pruning가 종종 불규칙한 아키텍처를 만들어내는 것과 달리, 레이어 가지치기는 Transformer의 전체 층(즉, Attention과 MLP 모듈 모두)을 제거하므로, 배포 및 가속화가 더 용이하다.

- ShortGPT (층 입력과 출력 간의 cosine similarity를 사용해 각 층의 중요도를 평가하고, 가장 중요도가 낮은 층을 제거)

- SLEB (ppl + iterative! pruning)

- Shortened LLaMA (tayler, ppl (한번에구해둠) + LoRA)

- UIDL (각 층 간의 각도 거리(angular distance)를 도입하여 연속적인 층을 식별하고 제거하며, 이후 QLoRA 적용)

- LLM-Streamline (cosine sim + 연속적 레이어 선택 + lightweight layer로 대체)

 

 

 

3. Method

3.1. Preliminaries on LLM Layer Pruning

- transformer layer의 기본 식.

X: Input activation

theta : parameters

 

Pruning Metrics.

보통 cosine similarity[ShortGPT, LLMStremline], gradient-based score[Shortened Llama,LLM-Pruner] , perlexity-based score[Shortened Llama, SLEB]를 사용한다.

 

Layer Pruning.

- 프루닝 이후 식

l* 번째 레이어부터 n개의 연속된 레이어가 제거되었을 경우, l*의 인풋이 l*+n 번째 파라미터(레이어)에 들어간다.

 

그런데, 가지치기 경계(pruning interface)에서 channel magnitude의 큰 불일치를 유발하며, 이는 모델 성능을 심각하게 저하 시킨다는 것을 Figure 1에서 확인한다. (cont. sections 3.2 / 3.3)

 

 

3.2. Channel Magnitude Alignment

Layer-wise Channel Mismatch.

figure1(a)에 나타난 것처럼, hidden state의 크기는, layer와 channel에 따라 달라진다. 

* channel: 모델의 hidden dimension을 의미. Llama2-7b 기준으로 4096.

 

이를 완화하기 위해, channel-wise scaling factor를 statistically하게 계산한다. 

각 채널 k에 대해, calibration set을 사용하여, l*번째 레이어와 (l*+n)번째 레이어의 평균 activation magnitude의 비율을 계산한다.

이를 통해 scaling vector d 를 만들어낸다.

 

-> Channel wise 스케일링 진행

 

Quantitative Evaluation.

추가적인 scaling factor인 알파를 사용해서 d 주변에서 변형한다.

figure1(b)에 나타난 것처럼, 그냥 알파==1인 경우가 가장 잘 좋은 성능을 보였다. 여기서 벗어날 경우 성능 저하가 심하게 발생했다.

 

 

3.3. Token Magnitude Smoothing

Token-wise Scaling Mismatch

최신 연구에 따르면, [BOS]토큰이나 구분자토큰과 같은 특정 토큰에 대해 크기가 10^3이상인 거대한 outlier가 존재한다.

따라서 single channel scaling d_k만으로는 채널 내의 모든 토큰에 적합하지 않을 수 있다. (figure2(a))

X_i,k : the activations of channel k for batch i ( i번째 배치에 대한 채널k의 activation)

σ(·) : standard deviation

σ_d가 작을수록 토큰 간 스케일링이 일관됨을 의미한다. (표준편차가 작다는 것이므로)

그러나 LLaMA-2-7B에서 9개 레이어를 pruning할 때 σ_d= 2137.75로 나타나, 토큰 수준에서 심각한 불일치가 존재함을 보여준다.

 

Hadamard Transformation

-> 토큰별 scaling을 진행

최근 연구 [30, 34, 4, 45]에 따르면, Hadamard transform을 적용하면 outlier를 억제할 수 있다.

* Hadamard transform : 선형변환 방법. 모든 원소기 +1 또는 -1이고, orthogonal(직교)하다. 

 

 

 

1️⃣ Walsh–Hadamard 행렬 H₂ 만들기 (2x2)

- 1/root2는 정규화 상수. 이걸 곱해줘야 변환 후에도 벡터 길이가 바뀌지 않는다.

--> H_2를 벡터에 곱하면, 벡터를 45도 회전시키고, 반대로 뒤집은 성분까지 포함시킨다. 

 

2️⃣ 더 큰 행렬 H₂ⁿ 만들기 (재귀)

여기서 ⊗는 크로네커 곱(Kronecker product)

 

3️⃣ C가 2ⁿ이 아닌 경우

C = 2^n m \quad \Rightarrow \quad H_C = H_{2^n} \otimes H_m
• 만약 채널 수가 2의 제곱수가 아니면, 가장 큰 2의 제곱수 부분과 나머지를 나눠서 만들 수 있어요.
• 이렇게 해도 직교 성질은 그대로 유지됩니다.

 

 

Hadamard matrix의 직교성(H^T * H = I) 덕분에 다음 변환이 동일하게 적용된다:

** activation X에다가 H를 곱하면 activation 값들이 채널에 골고루 섞이고

** H^T를 다시 곱하면 원래 값으로 돌아온다.

즉, 정보 손실 없이 rotated activation을 진행한 것.

 

 

이 회전은 outlier를 모든 channel에 재분배하고, 채널 간 activation의 분포를 보다 균형 있게 만든다.
회전된 activation을 사용하면 모든 토큰에 동일한 스케일링 파라미터 d를 적용하기가 용이해지며, σ_d는 230.32까지 낮아진다.

(근데.. 아. (l*번째 레이어의) Output으로 받는 엑티베이션마다 저 H 행렬을 곱하고, 스케일링 파라미터 d를 곱한 것을 다음 레이어(l*+n)의 input으로 넣어준다고..? - 그럼 이건 (l*+n)번째 레이어의 원래 input과는 크게 상관없고, 얘의 아웃풋과 유사한 magnitude로 변환해서 넣어주는 거 같은데 어떤 의미가 있는건지??? 오호 그냥 이 변환이 전부네?)

 

 

3.4. LinearPatch: the Ultimate Recipe

 

먼저 X에다가 Hadamard transform을 적용한 뒤, 회전된 공간에서 D로 스케일링한다.

위 두 연산은 하나의 대칭행렬 P로 통합된다. 

마지막 등식은 스펙트럼 정리 [21]에서 유래하며, 즉 모든 실수 대칭 행렬은 직교 행렬(H)과 대각 행렬(D)로 분해할 수 있다는 것을 의미한다. (하..뭐라노..)

 

figure3. patch matrix P가 레이어가 제거된 LLM에서 발생하는 격차를 효과적으로 보완함을 보여준다. 또한 LINEARPATCH는 변환 오버헤드를 줄이고 효율적인 파인튜닝을 가능하게 하는데, 행렬 곱을 위한 단일 GEMM (General Matrix Multiplication_ 걍 일반적인 행렬 곱셈을 의미) 연산만 필요하며, 세 개의 별도 GEMM 연산이 필요하지 않다.

 

 

Memory-Efficient Offline Knowledge Distillation (학습!!!!!!!!!!!!!!!!!!!)

기존의 KD 방법은 Teacher과 Student을 모두 GPU 메모리에 올려야 하므로, LLM에서는 막대한 메모리 사용으로 인해 현실적으로 어렵다. 반면, 식 (9)를 기반으로 하는 LINEARPATCH는 메모리 효율적인 오프라인 증류 전략을 지원 : 티처모델의 입출력만 저장하고, distillation 과정동안은 offline으로 유지한다.

 

작은 training corpus X (예.5000개)개를 사용해서, 티처모델의 top-K개의 아웃풋 logit probability distribution o_t 와 그 인덱스를 추출한다. 실제로는 K=100으로 설정하여 전체 32K 어휘를 저장하는 것과 비교해 메모리 사용량을 320배 절감한다.

 

마찬가지로 학생모델에서 동일한 인덱스를 사용하여 top-K개의 아웃풋 logit probability distribution o_s를 수집한다.

 

이렇게 얻은 두 개의 logit probability distribution의 KL divergencee를 최소화하는 방향으로 patch matrix P를 최적화하는 방향으로 학습을 진행한다. 

 

파인튜닝 과정에서는 P에 대한 양의 정부호 제약(positive-definite constraint)을 제거하여 더 큰 유연성을 부여하고, 나머지 모델 파라미터는 freeze하여 메모리 사용을 최소화한다.

* positive-definite constraint : 행렬의 모든 고유값이 0 이상이어야 한다는 제약. 모델 연산 안정성 때문에 사용하는데, finetuning 자유도가 제한됨. 

이 과정 전체는 가벼워서, 예를 들어 LLaMA-2-7B 파인튜닝은 단일 NVIDIA V100 GPU에서 30분 만에 완료된다.

 

 

 

 

- 많아질수록 평균 성능이 증가하기는 하나, benefits과 costs 측면에서 k=100이 적합.

 

- MSE는 KL보다 안좋았고 오버피팅되는 현상이 나타났다고 함.

 

 

 

4. Experiments

4.1. Setup

Models and Baselines.

[Models]

LLaMA2-7b, 13b

LLaMA3-8b

Baichuan2-7b

DeepSeek-R1-Distill

 

[Baselines]

(gradient based)

LLM-Pruner

(ppl based)

SLEB

(Taylor based)

shortend Llama

(cosine sim based)

ShortGPT

LLM-Streamline

 

 

Evaluation.

(ppl)

WikiText-2

C4

PTB

(NLU)

MMLU

(QA)

ARC-e / c

BoolQ

BellaSwag

PIQA

WinoGrande

WSC273

Race-h

CoPA

 

- MMLU는 오피셜 코드 사용, 이외에는 lm-eval-harness 사용

 

4.2. Implementation Details

Calibration and Fine-tuning

 Calibration :

프루닝할 레이어를 정하고 channel-wise scaling 파라미터를 초기화하기 위한 calibration dataset이 필요함.

WikiText-2에서 sequence length 2048인 128개 데이터를 랜덤 샘플링함.

 

Appendix C

- wiki-2 데이터의 개수를 다르게하여 테스트함.

- 데이터양과 성능향상을 재봤을 때 128이 최적이다.

Appendix D

- 타겟도메인과 같은 도메인의 calibration set을 사용하면 성능이 향상됨.

- 도메인이 달라도 ppl은 거의 변하지 않음 -> 우리 방법의 안정성 (.........와 이렇게 해석하는구나...)

- 데이터 품질이 달라도 ppl은 거의 유사함 -> 우리 방법의 안정성 (22)

 

 

 For fine-tuning :
LINEARPATCH, we use AdamW with a learning rate of 1e−4, training for one epoch on 5,000 WikiText-2 sentences of length 2048

Appendix E

- 데이터양과 성능 향상폭을 비교했을 때, 5000이 최적의 값이다. 

 

Resource Consumption

- PyTorch 사용

- single NVIDIA V100 GPU with 24GB memory

- 7b 모델에서, LinearPatch의 초기화는 30초, fine-tuning은 30분만에 완료됨.

 

 

Pruning Configurations

 이전 연구를 따라 가지치기 비율을 30% 미만으로 제한

 

4.3. Main Results

 

우리는 먼저 LINEARPATCH의 학습 없는(training-free) 환경에서의 효과를 평가한다. 구체적으로, 프루닝된 대형 언어 모델(LLM)의 추론 능력과 언어 모델링 능력 유지 정도를 측정하는 데 널리 사용되는 상식 기반 질문 응답(QA) 벤치마크와 퍼플렉서티(PPL) 벤치마크에 초점을 맞춘다. 비교의 공정성을 위해, 고려된 모든 접근법은 fine-tuning을 수행하지 않는다. 특히, LLM-Pruner의 경우 LoRA 기반 파인튜닝 단계를 제외하며, LLM-Streamline은 공식 프로토콜을 따르되, 레이어 교체(layer replacement)와 오프라인 증류(offline distillation)를 제거한 변형을 LLM-Streamline (None)으로 표기한다. 추가 LLM 백본 및 벤치마크 결과는 부록 I(Appendix I)에서 확인할 수 있다.

(??????????????? LLM-stremline의 핵심 방법을 제거해버리면 어캄?????? 내가 저자였으면 극대노했을듯)

 

4.3.1. Comparison on Training-free Methods

 

Results on QA Benchmarks

 

Results on PPL Benchmarks

 

Results on PPL Benchmarks

 

 

4.3.2. Comparison on Post-training Methods

 

Results on QA Benchmarks

 

Results on PPL Benchmarks

 

 

4.4. Discussions and Ablation Studies

Tunable Parameters and Loss Functions

 

 

The Ingredients of LinearPatch

 

Online Inference Overhead

 

Offline Storage Overhead

 

 

5. Conclusion

 

 

6. Limitation and Broader Impact

Limitation

레이어 프루닝은 서로 다른 작업(task)에서 모델 성능에 불균형적인 저하를 초래할 수 있다. 예를 들어, 일부 질문 응답(QA) 작업은 여전히 강인하게 유지될 수 있지만, 복잡한 추론(complex reasoning)이나 문맥 의존(context-dependent) 작업은 성능이 크게 저하될 수 있다. 향후 연구에서는 효율성 향상과 작업별 성능 간의 균형(trade-off)을 평가할 수 있는 체계를 구축할 필요가 있다.

 

Broader Impact

레이어 프루닝 방법은 LLM의 배포에 필요한 계산 비용을 크게 줄여, 더 많은 사용자에게 접근성을 높인다. 그러나 이러한 방법은 LLM에 내재된 social biases을 해결하지 못하며, 이러한 편향은 종종 training data에서 비롯되어 공정성과 포용성(fairness and inclusivity)에 영향을 줄 수 있다. 따라서 LLM을 윤리적으로 배포하는 것이 매우 중요하다.