그라디언트 트레이스 기반 제로샷 LLM 프루닝

초록

TraceNAS는 사전 학습된 대형 언어 모델의 손실 지형 정렬을 측정하는 스케일‑불변 제로샷 프록시를 이용해, 깊이와 폭을 동시에 구조적 프루닝하는 훈련‑무료 NAS 프레임워크이다. 8.5시간 내 단일 GPU에서 탐색을 마쳐, 기존 훈련‑인식 방식 대비 10배 적은 GPU‑시간으로 Llama·Qwen 계열 모델을 경쟁력 있게 압축한다.

상세 분석

본 논문은 대형 언어 모델(LLM)의 구조적 프루닝을 두 차원(깊이와 폭)에서 동시에 최적화하는 새로운 NAS 접근법인 TraceNAS를 제안한다. 기존의 레이어·헤드·채널 단위 중요도 평가 방식은 각 부분을 독립적으로 판단해 전역적인 구조적 의존성을 반영하지 못한다는 한계가 있다. 반면, 훈련‑인식(pruning‑aware) 방법은 전체 모델의 손실 지형을 고려하지만, 후보 모델마다 미세 조정과 대규모 토큰 샘플링을 필요로 하여 탐색 비용이 포스트‑프루닝 재학습과 맞먹는 수준으로 비싸다.

TraceNAS는 “기능적 상속”(functional inheritance)이라는 개념을 도입한다. 사전 학습된 모델이 손실 지형의 넓고 평탄한 영역에 위치한다는 가정 하에, 프루닝 후 모델이 원본 모델과 얼마나 동일한 방향의 그래디언트 흐름을 유지하는지를 측정한다. 이를 위해 저차원 서브스페이스에서 그래디언트 트레이스를 추출하고, 프루닝 전·후 모델의 그래디언트 트레이스 간 피어슨 상관계수를 가중합한 스케일‑불변 지표 Φ를 정의한다. 피어슨 상관계수는 절대 크기 차이를 무시하고 방향성만을 평가하므로, 구조적 프루닝이 초래하는 활성화·그래디언트 스케일 변동에 강인하다.

검색 공간은 각 트랜스포머 블록에 대해 (1) 블록을 유지하거나 스킵하는 깊이 마스크 d와 (2) 어텐션·MLP 서브블록별 파라미터 유지 비율 κ를 조합한 이산·연속 혼합 형태로 설계된다. κ는 활성화‑가중치 곱 I(κ)·j 를 이용해 채널별 중요도를 산출하고, 이를 이진 마스크로 변환해 인‑플레이(mask‑in‑place) 방식으로 가중치를 일시적으로 수정한다. 이렇게 하면 후보 모델을 실제로 복제하지 않아도 한 번의 포워드·백워드 연산으로 모든 후보의 그래디언트 트레이스를 수집할 수 있어 메모리와 연산량을 크게 절감한다.

진화적 탐색(evolutionary search)에서는 초기 인구를 무작위 생성한 뒤, 교차·돌연변이 연산을 통해 새로운 깊이·폭 조합을 만든다. 각 후보는 Φ 값을 통해 순위가 매겨지며, 상위 후보가 다음 세대로 전달된다. 탐색 과정 전체는 단일 GPU(NVIDIA A100)에서 8.5시간 만에 수백 개의 후보를 평가한다.

실험에서는 Llama‑7B·13B·30B와 Qwen‑7B·14B 모델을 대상으로, CommonSenseQA, ARC‑Easy/Challenge, GSM‑8K 등 다양한 추론·수학 벤치마크에서 기존 훈련‑인식 프루닝 기법(예: ShearedLlama, DarwinLM, PUZZLE)과 비교했다. TraceNAS가 찾은 비균일 프루닝 구조는 파라미터 30%50% 감소에도 불구하고, 포스트‑프루닝 재학습 후 정확도 손실이 12% 이내에 머물렀으며, 전체 탐색 비용은 훈련‑인식 방식 대비 평균 10배 이상 절감되었다. 또한, Φ와 실제 다운스트림 성능 간의 스피어만 ρ=0.94, 켄달 τ=0.82라는 높은 상관성을 보이며, 제로샷 프록시가 실제 성능을 잘 예측함을 입증했다.

핵심 기여는 (1) 그래디언트 트레이스를 이용한 스케일‑불변 제로샷 프록시 설계, (2) 깊이·폭을 동시에 탐색하는 효율적인 검색 공간 인코딩, (3) 인‑플레이 마스킹을 통한 메모리 효율성, (4) 훈련‑무료 방식으로도 훈련‑인식 프루닝과 동등하거나 우수한 압축 성능을 달성한 점이다. 이 접근법은 앞으로 더욱 큰 규모의 LLM에 대한 구조적 압축을 저비용으로 수행할 수 있는 길을 열어준다.