💻 computer science

SHIFT-LLM: Distribution Shift Correction in Depth-Pruned LLMs

SHIFT-LLM은 트랜스포머 블록을 제거함으로써 발생하는 분포 변화를 교정하기 위해 폐쇄형 회귀(closed-form regression)를 통해 보정된 경량 선형 잔차 어댑터(Linear Residual Adapters)를 삽입함으로써, 깊이 프루닝(depth-pruned)된 거대 언어 모델의 정확도를 복원하는 학습이 필요 없는 포스트 프루닝 프레임워크이다.

원저자: Ali Bahri, Hang Li, Hongliang Li, Zhitang Chen

게시일 2026-08-27

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Ali Bahri, Hang Li, Hongliang Li, Zhitang Chen

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

거대 언어 모델(LLM)은 오늘날 가장 진보된 인공지능 도구들의 핵심 엔진으로, 이야기를 쓰고, 문제를 해결하며, 복잡한 질문에 답할 수 있는 능력을 갖추고 있습니다. 이러한 시스템은 마치 다층 건물처럼 디지털 처리 장치들이 층층이 쌓인 구조로 구축되어 있으며, 각 층은 아래층에서 전달된 정보를 정교하게 다듬습니다. 건물의 층수가 많을수록 최종 출력물은 더 상세하고 정교해질 수 있지만, 이는 동시에 구조를 매우 무겁고 운영 비용이 많이 들게 만듭니다. 전력이나 메모리가 제한된 기기에서의 많은 실용적인 응용 분야에서는 이 거대한 모델들을 사용하기에 너무 큽니다. 모델을 더 작게 만들기 위해 연구자들은 건물에서 층 전체를 제거하는 방법인 '깊이 프루닝(depth pruning)'을 개발했습니다. 이 과정은 하중을 줄이고 속도를 높이는 데는 성공적이지만, 종종 남은 층들이 비틀거리게 만듭니다. 제거된 층들은 위층으로 전달될 특정 유형의 정보를 전달하도록 설계되었기 때문에, 이를 제거하면 상위 레이어들이 혼란을 겪으며 더 이상 학습된 내용과 일치하지 않는 신호를 받게 됩니다. '분포 변화(distribution shift)'라고 알려진 이 불일치는 모델의 정확도를 떨어뜨리고 실수를 유발하며, 개발자들이 오류를 수정하기 위해 상당한 시간과 컴퓨팅 자원을 들여 모델을 재학습시켜야 하는 상황을 초래합니다.

화웨이 노아의 방주 연구소(Huawei Noah's Ark Lab)의 연구팀은 재학습 없이 이 문제를 해결하는 SHIFT-LLM이라는 새로운 접근 방식을 도입했습니다. 누락된 층을 다시 구축하거나 남은 층을 다시 가르치는 대신, 이들의 방법은 레이어가 제거된 모든 지점에 작고 가벼운 보정 모듈을 삽입합니다. 만약 건물에서 한 층을 제거했다면, 그 빈자리를 완벽하게 연결하는 맞춤형 경사로를 설치하여 상층부 사람들이 마치 사라진 층이 여전히 존재하는 것처럼 정확한 지침을 받을 수 있도록 한다고 상상해 보십시오. 디지털 세계에서 이 경사로는 정보의 직접적인 경로를 보존하면서도 미세하고 계산된 조정을 더하는 선형 어댑터(linear adapter)입니다. 이 조정은 제거된 레이어가 기여했을 구체적인 역할을 모방하도록 설계되었으며, 결과적으로 나머지 모델이 사라진 층이 여전히 존재하는 것처럼 속이는 효과를 냅니다.

이 방법의 탁월함은 그 조정 값을 계산하는 방식에 있습니다. 연구진은 누락된 레이어가 무엇을 했어야 하는지 알아내기 위해 수천 시간의 훈련을 돌리는 대신, 소량의 샘да 데이터를 사용하여 정확히 어떤 정보가 손실되었는지 측정합니다. 그런 다음 단순한 수학적 계산을 통해 손실된 정보를 복구하는 데 필요한 정밀한 보정치를 결정합니다. 이 과정은 완전히 자동화되어 있으며, 일반적으로 고장 난 모델을 수정하는 데 필요한 복잡하고 반복적인 학습 과정인 '경사 기반 최 optimization(gradient-based optimization)'을 필요로 하지 않습니다. 이 폐쇄형 계산(closed-form calculation)을 사용함으로써, 연구팀은 단 몇 백 개의 예시만을 사용하여 단 몇 분 만에 모델의 내부 상태를 교정할 수 있으며, 이는 보통 며칠의 컴퓨팅 시간이 걸리는 작업입니다. 그 결과, 프루닝된 모델은 적은 수의 레이어를 가진 모델의 속도와 효율성을 유지하면서도 원래의 거대한 모델과 거의 동일한 정확도를 유지하게 됩니다.

연구진은 실험을 통해 이 기술을 15억 파라미터 규모의 작은 모델부터 140억 파라미터 규모의 큰 모델까지 포함하는 다섯 가지 서로 다른 LLM 제품군에 대해 테스트했습니다. 그들은 어떤 레이어를 제거할지 결정하기 위해 여섯 가지 서로 다른 방법을 적용했으며, 일반적인 지식과 추론을 테스트하기 위해 설계된 일곱 가지 벤치마크에서 결과를 평가했습니다. 결과는 일관적이었습니다. 보정 모듈은 거의 모든 구성에서 프루닝 중에 손실된 정확도를 성공적으로 회복했습니다. 특히 80억 파라미터 모델의 경우, 이 방법은 표준 벤치마크에서 무려 15.7포인트의 정확도를 회복했는데, 이는 통상적으로 광범-하고 값비싼 재학습이 필요한 수준의 향상입니다. 모델이 프루닝 후 이미 미세 조정(fine-tuning)되고 있는 상황에서도 이 보정 모듈을 추가하면 추가적인 성능 향상을 얻을 수 있었으며, 이는 두 방식이 성능을 더욱 높이기 위해 잘 결합될 수 있음을 시사합니다.

또한 연구진은 이러한 보정 모듈을 공간 절약을 위해 더 압축할 수 있으며, 여러 레이어가 연속으로 제거될 때 이를 하나로 병합할 수 있다는 점을 입증하여, 프루닝의 효율성 이점이 보정 작업의 오버헤드로 인해 상쇄되지 않도록 했습니다. 이 연구는 거대 언어 모델을 작게 만드는 데 있어 주요 장애물이 단순히 부품을 제거하는 것이 아니라, 그들 사이의 정보 흐름을 방해하는 것임을 보여줍니다. 단순하고 표적화된 조정을 통해 그 흐름을 복구하는 데 집중함으로써, SHIFT-LLM은 강력한 AI 모델을 일상적인 용도로 사용할 수 있게 만드는 일반적이고 효율적인 방법을 제시합니다. 이 연구는 적절한 보정이 있다면, 우리는 명료하게 사고하는 능력을 잃지 않으면서도 디지털 뇌의 무겁고 중복된 부분을 덜어낼 수 있음을 확인시켜 주며, 더 빠르고 저렴하며 접근하기 쉬운 인공지능의 문을 열어줍니다.

기술 요약: SHIFT-LLM

문제 정의

LLM(대규모 언어 모델)의 깊이 프루닝(Depth pruning)은 전체 트랜스포머 블록을 제거하여 추론 효율성을 높이는 것을 목표로 합니다. 그러나 이 과정은 결정적인 문제인 **분포 변화(distribution shift)**를 야기합니다. 중간 블록이 제거되면, 해당 블록이 생성했을 잔차 업데이트(residual update)가 손실됩니다. 결과적으로 다운스트림 레이어들은 사전 학습 시 경험했던 것과는 체계적으로 다른 은닉 상태(hidden states)를 전달받게 됩니다. 이러한 변화는 네트워크를 통해 전파되며, 심각한 정확도 저하를 초러합니다. 프루닝 후 미세 조정(post-pruning fine-tuning)을 통해 성능을 회복할 수 있지만, 이는 상당한 계산 자원, 최적화 시간, 그리고 그래디언트 기반 학습을 필요로 하므로 자원이 제한된 시나리오에서는 적용이 어렵습니다.

방법론: SHIFT-LLM

저자들은 그래디언트 계산 없이 분포 변화를 완화하기 위해 설계된, 학습이 필요 없는(training-free) 포스트 프루닝 교정 프레임워크인 SHIFT-LLM을 제안합니다. 핵심 메커니즘은 각 프루닝 지점에 삽입되는 **선형 잔차 어댑터(Linear Residual Adapter, LRA)**입니다.

1. 선형 잔차 어댑터 (LRA) 공식화

LRA는 원래의 출력 은닉 상태를 처음부터 재구성하는 대신, 원래의 **항등 경로(identity pathway)**를 보존하면서 가벼운 **아핀 잔차 교정(affine residual correction)**을 추가합니다.

구조: 프루닝된 레이어 $\ell$ 에 대해, 원래의 출력 $h^{(\ell)} = h^{(\ell-1)} + f_\ell(h^{(\ell-1)})$ 은 $\hat{h}^{(\ell} = h^{(\ell-1)} + g_\ell(h^{(\ell-1)})$ 로 근사됩니다.
교정 함수: 어댑터 $g_\ell(\cdot)$ 는 아핀 변환으로 모델링됩니다: $g_\ell(h) = hA^{(\ell)} + \mathbf{1}(b^{(\ell)})^\top$ , 여기서 $A^{(\ell)}$ 은 가중치 행렬이고 $b^{(\ell)}$ 은 편향 벡터입니다.
목표: LRA는 전체 은닉 상태를 근사하는 것이 아니라, 누락된 잔차 업데이트( $\delta^{(\ell)} = f_\ell(h^{(\ell-1)})$ )를 구체적으로 근사합니다. 이는 모델이 사전 학습된 잔차 구조를 활용함으로써 근사 문제를 단순화합니다.

2. 폐쇄형 추정 (학습 불필요)

LRA의 파라미터( $A^{(\ell)}$ 및 $b^{(\ell)}$ )는 역전파나 반복적 최적화 없이 다음과 같이 추정됩니다:

데이터 수집: 원래의 프루닝되지 않은 모델이 소량의 홀드아웃 교정 데이터셋(held-out calibration set)에 대해 순방향 패스(forward pass)를 수행합니다. 입력 은닉 상태( $H^{(\ell)}$ )와 그에 대응하는 실제 잔차 업데이트( $\Delta^{(\ell)}$ )를 수집합니다.
릿지 회귀(Ridge Regression): 예측된 잔차 업데이트와 실제 잔차 업데이트 사이의 오차를 최소화하기 위해 폐쇄형 최소제곱 회귀(ridge regression)를 통해 파라미터를 해결합니다:
$X^{(\ell)*} = \left((H^{(\ell)}_{aug})^\top H^{(\ell)}_{aug} + \lambda I\right)^{-1} (H^{(\ell)}_{aug})^\top \Delta^{(\ell)}$
여기서 $H^{(\ell)}_{aug}$ 는 편향 컬럼을 포함하며, $\lambda$ 는 규제 계수입니다.
효율성: 이 과정은 각 프루닝된 레이어당 단 한 번의 순방향 패스와 선형 연산(linear solve)만을 필요로 하며, 그래디언트 계산이나 미세 조정을 피할 수 있습니다.

3. 실용적 확장

저계수 인수분해(Low-Rank Factorization): 적합된 가중치 행렬 $A^{(\ell)}$ 은 절단된 특이값 분해(truncated SVD)를 통해 파라미터 수를 $d^2$ 에서 $2dr + d$로 압축할 수 있습니다.
정확한 병합(Exact Merging): LRA는 아핀 사상(affine mapping)을 정의하므로, 인접한 프루닝 지점에 있는 연속적인 LRA들은 근사 오차 없이 하나의 변환으로 수학적으로 병합될 수 있으며, 이를 통해 압축 효율을 더욱 높일 수 있습니다.

주요 기여

SHIFT-LLM 프레임워크: 깊이 프루닝으로 인한 은닉 상태의 분포 변화를 명시적으로 겨냥하는 학습 불필요 포스트 프루닝 교정 방법입니다.
선형 잔차 어댑터 (LRA): 항등 경로를 보존하면서 폐쇄형 회귀를 통해 교정된 가벼운 아핀 교정을 사용하여 누락된 잔차 업데이트를 근사하는 새로운 구성 요소입니다.
모듈성 및 일반성: 이 프레임워크는 특정 레이어 선택 전략(예: Block Influence, Taylor sensitivity, Magnitude)에 구애받지 않으며 다양한 모델 제품군에 적용 가능합니다.
상호 보완성: 이 방법은 기존의 매개변수 효율적 미세 조정(예: LoRA)과 함께 작동하도록 설계되어, 미세 조정만으로 얻을 수 있는 성능 이상의 회복력을 제공합니다.

실험 결과

저자들은 5개의 모델 제품군(Qwen2-1.5B, Qwen1.5-7B, Llama-3.1-8B-Instruct, Vicuna-7B, Qwen2.5-14B), 6개의 레이어 선택 기준, 그리고 7개의 제로샷 벤치마크를 통해 SHIFT-LLM을 평가했습니다.

학습 불필요 회복(Training-Free Recovery): 미세 조정을 전혀 수행하지 않고도, SHIFT-LLM은 깊이 프루닝으로 인해 손실된 정확도를 일관되게 회복했습니다.
- Llama-3.1-8B-Instruct에서 역순 프루닝(Reverse-order pruning, 25% 깊이)을 적용했을 때, 이 방법은 평균 제로샷 정확도에서 +15.74 포인트의 이득을 달성했습니다.
- Qwen1.5-7B의 경우, Magnitude- $\ell_1$ 프루닝 하에서 +6.88 포인트의 이득을 기록했습니다.
- WikiText-2에 대한 퍼플렉시티(Perplexity)는 대부분의 설정에서 크게 감소했습니다 (예: Llama-3.1-8B 역순 프루닝 하에서 1000 포인트 이상의 퍼플렉시티 감소).
확장성: 이 방법은 Qwen2.5-14B로 확장해도 효과가 유지되었으며, 제로샷 정확도에서 최대 +5.36 포인트의 이득을 보였습니다.
미세 조정과의 호환성: LoRA 또는 부분 레이어 미세 조정과 결합했을 때, SHIFT-LLM은 상호 보완적인 회복을 제공하여 대부분의 평가된 설정에서 성능을 더욱 향상시켰습니다 (예: Llama-3.1-8B-Instruct에서 LoRA와 결합 시 +1.14 포인트 향상).
SOTA 비교: Navigation LLM 파이프라인과 통합되었을 때, SHIFT-LLM은 Llama-3.1-8B-Instruct의 평균 정확도를 59.08에서 59.54로 개선하여, 강력한 베이스라인을 저하시키지 않으면서도 가치를 더한다는 것을 입증했습니다.

의의 및 주장

본 논문은 SHIFT-LLM이 깊이 프루닝에 내재된 분포 변화 문제에 대해 단순하고, 일반적이며, 효과적인 솔루션을 제공한다고 주장합니다. 원래 모델의 잔차 구조를 보존하면서 누락된 잔차 업데이트를 명시적으로 근사함으로써, 이 방법은 재학습의 계산 비용 없이 분포 변화를 완화합니다.

저자들은 이 접근 방식이 특히 포스트 프루닝 학습이 바람직하지 않거나 불가능한 환경에서 매우 유의미하다고 강조합니다. 단 몇 백 개의 교정 샘플만을 사용하고 그래디언트 계산 없이 상당한 정확도 회복(최대 +15.7 포인트)을 달라는 점은 이 프레임워크를 광범위한 미세 조정의 실질적인 대안으로 자리매김하게 합니다. 또한, 연속적인 어댑터를 병합하고 저계수 인수분해를 통해 가중치를 압축할 수 있는 능력은 깊이 프루닝의 효율성 이점을 유지하거나 오히려 강화할 수 있음을 보장합니다.