💬 NLP

TraceNAS: Zero-shot LLM Pruning via Gradient Trace Correlation

TraceNAS는 그래디언트 트레이스 상관관계(gradient trace correlation)를 활용하여 원래의 사전 학습된 모델과 손실 지형(loss landscape)을 정렬함으로써 최적의 비균등 프루닝된 거대 언어 모델을 식별하는 매우 효율적인 훈련 불필요 신경망 구조 탐색(Neural Architecture Search) 프레임워크로, 훈련 인지 방식(training-aware methods)의 성능에 필적하는 성과를 훨씬 적은 계산 비용으로 달성합니다.

원저자: Prajna G. Malettira, Manish Nagaraj, Arjun Roy, Shubham Negi, Kaushik Roy

게시일 2026-02-04

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Prajna G. Malettira, Manish Nagaraj, Arjun Roy, Shubham Negi, Kaushik Roy

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

당신에게 인류 지식의 총합이 담긴 거대하고 믿을 수 없을 정도로 똑똑한 도서관(거대 언어 모델, LLM)이 있다고 상상해 보세요. 이 도서관은 너무 크고 무거워서 들고 다니거나 일반적인 휴대폰에서 사용하는 것이 불가능합니다. 당신은 이 모델의 이야기하는 능력, 수수께끼를 푸는 능력, 혹은 농담을 이해하는 능력을 잃지 않으면서도 이를 주머니에 들어갈 만한 크기로 줄이고 싶습니다.

이것이 바로 TraceNAS라는 논문이 해결하고자 하는 문제입니다.

문제점: 잘못된 것을 잘라내는 것

보통 사람들이 이 거대한 모델들을 축소하려고 할 때, 그들은 서툰 정원사처럼 행동합니다. 그들은 "모든 나무에서 가지를 50%씩 잘라내자"라거나 "가장 무거운 잎사귀들을 제거하자"라고 말할 수 있습니다. 이것을 **균등 가지치기(uniform pruning)**라고 부릅니다.

하지만 여기에는 함정이 있습니다: 모든 뇌의 부분(혹은 도서관의 부분)이 다 같은 것은 아닙니다. 어떤 부분은 복잡한 논리를 담고 있는 "중요 장기"인 반면, 다른 부분은 쉽게 쳐낼 수 있는 "지방"일 뿐입니다. 만약 잘못된 중요 장기를 잘라낸다면, 모델은 모든 것을 잊어버릴 것입니다. 반대로 지방만 잘라낸다면, 모델은 여전히 너무 무겁게 남을 것입니다.

기존의 방법들은 다음 두 가지 방식으로 무엇을 잘라낼지 결정하려고 시도합니다:

한 번에 한 부분씩 살펴보기: 특정 뉴런이 중요한지 확인하지만, 그 뉴런이 나머지 뇌와 어떻게 소통하는지는 놓칩니다.
모델을 자르는 동안 다시 학습시키기: 모델을 축소하는 동시에 모델에게 생각하는 법을 다시 가르치려고 합니다. 이것은 자동차를 운전하는 법을 배우는 동시에 엔진을 재건축하는 것과 같습니다. 시간이 엄청나게 오래 걸리고 막대한 양의 연료(컴퓨팅 파워)가 필요합니다.

해결책: TraceNAS ( "그래디언트 흔적" 탐정)

저자들은 TraceNAS라고 불리는 새로운 방법을 제안합니다. 이것은 모델을 직접 건드리거나 다시 가르치지 않고도 모델의 "영혼"을 들여다볼 수 있는 첨단 X-레이 스캐너라고 생각하면 됩니다.

작동 방식은 다음과 같습니다 (쉬운 비유를 사용하겠습니다):

1. 원래 마음의 "메아리"

원래의 거대한 모델이 완벽한 요리를 만든 마스터 셰프라고 상상해 보세요. 당신이 음식을 한 입 먹었을 때, 그 맛은 혀에 특정한 "흔적"을 남깁니다.

기존 방식: 새로운 작은 레시피가 효과가 있는지 보려면, 요리를 전체 다 만들어 맛을 본 뒤, 맛이 없다면 처음부터 다시 시작해야 합니다.
TraceNAS 방식: 전체 요리를 만드는 대신, TraceNAS는 **재료와 레시피 단계(그래디언트)**를 보고 맛이 제대로 나올지 예측합니다. 그것은 작게 축소된 버전의 "맛의 흔적"이 원래 마스터 셰프의 흔적과 일치하는지 확인합니다.

2. "그림자" 테스트 (그래디언트 흔적 상관관계)

이 논문은 **그래디언트 흔적 상관관계(Gradient Trace Correlation)**라는 수학적 개념을 사용합니다.

원래 모델을 잔잔한 바다를 항해하는 거대한 배라고 상상해 보세요. 배는 뒤에 특정한 항적(파도의 흔적)을 남깁니다.
당신이 더 작은 배(가지치기 된 모델)를 만들려고 할 때, TraceNAS는 이렇게 묻습니다: "이 작은 배가 큰 배의 항적과 똑같이 생긴 항적을 남기는가?"
만약 항적이 일치한다면, 그것은 작은 배가 같은 경로를 따르고 있으며, 약간의 연습만 거치면 목적지에 잘 도착할 것임을 의미합니다. 만약 항적이 혼란스럽다면, 그 배는 충돌할 것입니다.

3. "저차원(Low-Rank)" 지름길

거대한 배의 항적을 계산하는 데는 보통 슈퍼컴퓨터가 필요합니다. TraceNAS는 영리합니다. 배의 움직임이 주로 몇 가지 주요 방향으로 일어난다는 점을 깨달았습니다. 이들은 저차원(Low-Rank) 기술을 사용하여 가장 중요한 몇 가지 방향만을 살펴봅니다.

비유: 바다의 모든 잔물결을 측정하는 대신, 가장 큰 세 개의 파도만을 측정하는 것입니다. 이를 통해 그들은 단 하나의 그래픽 카드(GPU)만으로 단 8.5시간 만에 이 테스트를 수행할 수 있었습니다. 기존 방식들이 전체 컴퓨터 클러스터를 사용하여 며칠 또는 몇 주가 걸렸던 것과 대조적입니다.

결과: 빠르고, 저렴하며, 똑똑함

저자들은 이 방법을 Llama와 Qwen 같은 유명한 모델들에 테스트했습니다.

속도: 그들은 단 한 대의 컴퓨터로 8.5시간 만에 최적의 "절단" 지점을 찾아냈습니다. 기존 방식들은 10배 더 오래 걸렸고 10배 더 많은 에너지를 사용했습니다.
정확도: 결과로 나온 작은 모델들은 최적의 절단 지점을 찾기 위해 몇 주 동안 학습된 모델들과 거의 동일한 성능을 보여주었습니다.
비균등성: 모든 것을 똑같이 자르는 "서툰 정원사"와 달리, TraceNAS는 "맞춤형 피팅"을 찾아냈습니다. 복잡하고 무거운 뇌의 부분은 온전히 유지하면서 지방만을 쳐내어, 매우 똑똑하면서도 효율적인 모델을 만들어냈습니다.

핵심 요약

TraceNAS는 거대한 AI의 뇌를 다시 학습시키거나 막대한 전기 비용을 들일 필요 없이, 주머니 크기의 버전으로 축소할 수 있게 해주는 도구입니다. 이는 작은 뇌의 "그림자"가 큰 뇌의 "그림자"와 일치하는지 확인함으로써 가능해집니다. 그림자가 일치한다면, 그 작은 뇌는 준비가 된 것입니다.

이것은 강력한 AI를 거대한 데이터 센터 없이도 일반 기기에서 실행할 수 있게 해주며, 동시에 엄청난 시간과 에너지를 절약해 줍니다.

기술 요약: TraceNAS

문제 정의

대규모 언어 모델(LLM)의 배포는 방대한 파라미터 수와 계산 요구 사항에 의해 제약을 받습니다. 구조적 프루닝(Structured Pruning)은 건축적 단위(예: 어텐션 헤드, 트랜스포머 블록)를 제거함으로써 효율적인 배포를 위한 경로를 제공하지만, 기존 방식들은 다음과 같은 중대한 한계에 직면해 있습니다:

지역적 vs 전역적 의존성: 대부분의 프러닝 방법은 로컬 휴리스틱(예: 활성화 가중치 크기 또는 헤시안 기반 재구성)을 사용하여 레이어, 헤드 또는 채널의 중요도를 개별적으로 평가합니다. 이러한 접근 방식은 복잡한 전역적 구조 의존성과 프러닝으로 인해 발생하는 오류의 전파 과정을 무시합니다.
학습 인지 탐색(Training-Aware Search)의 계산 비용: 전역적 의존성을 포착하려는 시도들은 대개 학습 인지 신경 구조 탐색(NAS)에 의존합니다. 이러한 방식은 과적합을 방지하기 위해 반복적인 그래디언트 업데이트와 지속적인 미학습 토큰 스트림을 필요로 하며, 이로 인해 탐색 과정이 프러닝 후 회복 학습 자체만큼이나 많은 계산 비용을 소모하게 됩니다.
기존 제로샷 프록시(Zero-Shot Proxy)의 부적절성: 학습이 불필요한 기존 프록시들은 (그래디언트 안정성, 손실 함수 지형의 매끄러움 등) 학습 가능성에 초점을 맞추어 무작위로 초기화된 아키텍처를 식별하도록 설계되었습니다. 이들은 "기능적 상속(Functional Inheritance)"—즉, 프러닝된 모델이 사전 학습된 손실 지형을 보존함으로써 성능을 회복할 수 있는 능력—을 평가하는 데 실패합니다.

방법론

TraceNAS는 탐색 시점에 별도의 학습 없이 LLM의 깊이(Depth)와 너비(Width)를 공동 최적화하도록 설계된 학습 불필요(Training-free) NAS 프레임워크로 제안되었습니다.

1. 탐색 공간 및 인코딩

본 프레임워크는 사전 학습된 LLM을 슈퍼 네트워크(Super-network)로 취급합니다. 탐색 공간 $G$ 는 다음의 결합 분포로 정의됩니다:

깊이 ( $d$ ): $L$ 개의 트랜스포머 블록에 대한 이산 마스크(Discrete mask)로, 비균일한 깊이 축소를 허용하면서 잔차 연결(Residual stream)의 연속성을 유지합니다.
너비 ( $\kappa$ ): 어텐션 및 MLP 서브 블록에 대한 파라미터 유지 비율로, 모델 전체에 걸쳐 이질적인 희소성(Heterogeneous sparsity)을 가능하게 합니다.

2. 인플레이스(In-Place) 아키텍처 구현

수천 개의 후보 모델을 인스턴스화하는 메모리 오버헤드를 피하기 위해, TraceNAS는 인플레이스 마스킹 전략을 사용합니다:

너비 마스킹: 가중치 채널에 대한 이진 마스크를 생성하기 위해 활성화 가중치 곱(Wanda와 유사)을 사용합니다. 이 $O(d^2)$ 휴리스틱은 이상치 특징(Outlier features)을 보존하면서 기능적으로 중요한 채널을 식별합니다.
깊이 마스킹: 비활성화된 경우 전체 블록을 우회하며, 활성화 값을 잔차 연결을 통해 라우팅합니다.
실행: 프레임워크는 최소한의 캘리브레이션 세트에 대해 그래디언트를 계산하기 위해 베이스 모델의 가중치를 일시적으로 수정하며, 이후 원래의 가중치를 복구하여 다음 반복을 위해 베이스 모델을 보존합니다.

3. 제로샷 프록시: 그래디언트 트레이스 상관관계(Gradient Trace Correlation)

핵심 혁신은 프러닝된 후보( $M_{sub}$ )와 사전 학습된 베이스 모델( $M_{base}$ ) 사이의 **기능적 정렬(Functional Alignment)**을 측정하는 스케일 불변(Scale-invariant) 프록시 $\Phi$ 입니다.

그래디언트 트레이스: 캘리브레이션 세트에 대해 그래디언트가 계산됩니다. 확장성을 보장하기 위해, 이들은 LLM의 본질적인 저차원성을 활용하여 메모리 오버헤드를 수십 배 줄이는 저차원(LoRA) 서브스페이스로 투영됩니다.
정렬 메트릭: 프록시는 각 레이어에 대해 프러닝된 모델과 베이스 모델의 표준화된 그래디梯度 트레이스 간의 **피어슨 상관 계수(Pearson Correlation Coefficient)**를 계산합니다. 이는 스케일 불변성을 보장하며, 크기 변화보다는 방향성 정렬에 집중합니다.
희소성 가중 집계: 최종 점수 $\Phi$ 는 파라미터 유지 비율에 대응하는 레이어별 상관관계의 가중 합입니다. 이를 통해 글로벌 점수를 높은 용량의 영역에 고정시켜, 공격적으로 프러닝된 서브 블록에서 발생하는 노이즈가 순위를 왜곡하는 것을 방지합니다.

4. 진화적 탐색(Evolutionary Search)

TraceNAS는 이산 탐색 공간을 탐색하기 위해 진화 전략을 채택합니다:

초기화: 글로벌 중요도 사전 확률(Priors)을 기반으로 모집단이 생성됩니다.
진화: 모집단은 교차(Crossover)와 변이(Mutation)를 거칩니다. 깊이 진화는 비트 플립 변이와 연속 범위 교차를 사용하며, 너비 진화는 보간(Interpolation)과 가우시안 지터(Gaussian jitter)를 사용합니다.
선택: 후보들은 오직 프록시 $\Phi$ 에 의해 순위가 매겨지며, 프러닝 후 지속적 사전 학습(CPT) 동안 성능 회복 잠재력이 가장 높은 모델을 선택합니다.

주요 기여

확장 가능한 학습 불필요 NAS: TraceNAS는 탐색 시점의 학습 없이 비균일한 아키텍처를 발견할 수 있게 함으로써, 깊이와 너비의 공동 프러닝을 위한 통합 프레임워크를 제공합니다.
기능적 상속 프록시: 프러닝된 모델과 사전 학습된 모델 간의 그래디언트 트레이스 정렬을 평가하는 새로운 스케일 불변 그래디언트 기반 프록시를 도입했습니다. 이는 다운스트림 정확도와 높은 상관관계(Spearman $\rho = 0.94$ , Kendall $\tau = 0.82$ )를 달성합니다.
효율성: 저차원 그래디언트 서브스페이스와 인플레이스 마스킹을 활용하여 탐색 오버헤드를 줄임으로써, 단일 GPU에서 고충실도 모델 발견을 가능하게 합니다.

실험 결과

TraceNAS는 Llama-2-7B, Llama-3.1-8B, Qwen-2.5-14B를 대상으로 평가되었으며, 각각 2.7B, 4.6B, 8.4B 파라미터로 프러닝되었습니다.

탐색 효율성: TraceNAS는 A100 GPU 기준 8.5 GPU 시간 만에 최적의 아키텍처를 식별하였으며, 이는 DarwinLM과 같은 학습 인지 베이스라인 대비 10배 감소한 수치입니다. 또한 DarwinLM에 비해 10배 적은 98M의 탐색 토큰만을 사용했습니다.
성능:
- Llama-2-7B: TraceNAS는 8개의 추론 벤치마크에서 평균 62.81%의 최고 정확도를 달 기록하며, DarwinLM(62.57%)과 ShearedLLaMA(62.63%)를 능가했습니다. 특히, ShearedLLaMA는 TraceNAS가 사용한 것과 동일한 토큰 예산 내에서 성능이 크게 하락(60.86%)했습니다.
- Llama-3.1-8B & Qwen-2.5-14B: TraceNAS는 복잡한 추론 작업(예: ARC-Challenge, BoolQ)에서 균등 프러닝 베이스라인 및 학습 인지 방식(DarwinLM, E3-Pruner)보다 뛰어난 성능을 보였으며, 훨씬 적은 회복 토큰(Minitron의 94B 대비 20B)을 필요로 했습니다.
프록시 검증: 프록시는 다양한 하이퍼파라미터(샘플 크기, 컨텍스트 길이, LoRA 랭크)에 대해 높은 안정성을 보여, 제약된 환경에서도 그래디언트 트레이스 정렬이 모델 상속을 효과적으로 포착함을 확인했습니다.

의의 및 주장

본 논문은 TraceNAS가 구조적 프러닝의 핵심 병목 구간인 계산 비용 문제를 해결한다고 주장합니다. "학습 가능성"(무에서 유를 학습)에서 "상속"(사전 학습된 손실 지형 보존)으로 초점을 전환함으로써, TraceNAS는 최소한의 자원 소비로 고성능의 비균일 프러닝 모델을 발견할 수 있게 합니다.

저자들은 TraceNAS가 다음과 같은 강점을 가진다고 강조합니다:

탐색 시점의 학습 비용을 제거하여 대규모 모델 압축을 용이하게 만듭니다.
로컬 휴리스틱이 놓치는 전역적 구조 의존성을 포착하여 복잡한 추론 능력의 보존을 보장합니다.
아키텍처 탐색에 필요한 GPU 시간을 획기적으로 낮춤으로써 모델 압축의 탄소 발자국을 줄입니다.

이 연구는 TraceNAS가 밀집된 사전 학습 모델의 기능적 무결성을 유지하면서, 자원이 제한된 환경에서 효율적인 LLM을 배포하기 위한 확장 가능한 기초를 마련했음을 보여줍니다.