[PYTHON] 모델 가지치기(Pruning) 후 재학습(Fine-tuning) 성능 회복 방법과 3가지 핵심 차이 해결 전략

모델 가지치기(Pruning) 후 재학습(Fine-tuning) 성능

딥러닝 모델의 경량화 과정에서 발생하는 성능 저하 문제를 수학적, 실무적 관점에서 분석하고, 최적의 회복 전략을 제시합니다.

1. 모델 가지치기(Pruning)의 본질과 직면하는 과제

최신 딥러닝 모델은 수십억 개의 파라미터를 가지고 있어 모바일 기기나 엣지 컴퓨팅 환경에서 구동하기에 너무 무겁습니다. 모델 가지치기(Pruning)는 가중치 중 중요도가 낮은 것을 제거하여 파라미터 수를 줄이는 핵심 기술입니다. 하지만 가지치기 직후에는 모델의 정확도가 급격히 하락하며, 이를 원상복구 하기 위한 재학습(Fine-tuning) 과정은 단순한 학습보다 훨씬 정교한 전략을 필요로 합니다. 본 가이드에서는 단순히 가중치를 지우는 것을 넘어, 지워진 파라미터의 공백을 메우고 성능을 극대화하는 7가지 이상의 실무 예제와 수학적 정당성을 다룹니다.

2. 가지치기 기법에 따른 재학습 전략의 3가지 차이점

가지치기는 방식에 따라 재학습의 접근법이 완전히 달라집니다. 아래 표는 구조적 차이와 그에 따른 해결 방안을 요약한 것입니다.

구분 항목	비구조적 가지치기 (Unstructured)	구조적 가지치기 (Structured)
제거 대상	개별 가중치 (Weight-wise)	채널 또는 필터 (Channel/Filter-wise)
성능 저하 폭	비교적 낮음 (희소성 유지 시)	매우 높음 (네트워크 구조 변경)
재학습 핵심 방법	Weight Rewinding (LR 스케줄링)	Knowledge Distillation (지식 증류)
속도 개선 효과	전용 하드웨어 필요	CPU/GPU 즉시 속도 향상
해결 전략 수치	희소도 90% 이상에서도 회복 가능	필터 50% 제거 시 정밀 튜닝 필수

3. 성능 회복을 위한 수학적 정당성: Lottery Ticket Hypothesis

MIT의 Jonathan Frankle이 제안한 '복권 가설(Lottery Ticket Hypothesis)'에 따르면, 거대한 네트워크 안에는 독립적으로 훈련했을 때 원래 모델과 필적하는 성능을 내는 '운 좋은 부분망(Winning Ticket)'이 존재합니다. 가지치기 후 재학습이 성공하려면 이 부분망을 찾아내고, 가중치를 초기 학습 시점의 값으로 되돌리는(Weight Rewinding) 기법이 수학적으로 가장 안정적인 수렴을 보장합니다.

4. 실무 적용을 위한 Python 샘플 코드 (Example 7선)

PyTorch 프레임워크를 기준으로 현업에서 즉시 사용 가능한 단계별 가지치기 및 재학습 파이프라인입니다.

Example 1: 글로벌 비구조적 가지치기 적용

import torch
import torch.nn.utils.prune as prune

def apply_global_pruning(model, amount=0.3):
    # 모델 내의 모든 선형 레이어와 컨볼루션 레이어를 대상으로 전역 가지치기 수행
    parameters_to_prune = [
        (module, 'weight') for module in model.modules()
        if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.Linear)
    ]
    
    prune.global_unstructured(
        parameters_to_prune,
        pruning_method=prune.L1Unstructured,
        amount=amount,
    )
    print(f"Global pruning applied with amount: {amount}")

Example 2: 지식 증류(Knowledge Distillation) 기반 재학습

교사 모델(원본)의 출력을 학생 모델(가지치기 모델)이 따라가게 하여 성능 회복 속도를 2배 이상 높입니다.

import torch.nn.functional as F

def kd_loss(outputs, teacher_outputs, labels, T=3.0, alpha=0.5):
    # KL Divergence를 이용한 증류 손실 계산
    distillation_loss = F.kl_div(
        F.log_softmax(outputs / T, dim=1),
        F.softmax(teacher_outputs / T, dim=1),
        reduction='batchmean'
    ) * (T * T)
    
    student_loss = F.cross_entropy(outputs, labels)
    return alpha * student_loss + (1 - alpha) * distillation_loss

Example 3: 가중치 되감기(Weight Rewinding) 전략

def weight_rewinding(model, initial_state_dict, mask_dict):
    # 가지치기된 마스크를 유지하면서 가중치만 초기 값으로 복사
    current_state = model.state_dict()
    for name, param in current_state.items():
        if name in initial_state_dict:
            # 마스크가 1인 부분(살아남은 가중치)만 초기값으로 복구
            current_state[name].copy_(initial_state_dict[name] * mask_dict[name])
    model.load_state_dict(current_state)

Example 4: 학습률 스케줄러를 활용한 정밀 튜닝 (Cosine Annealing)

import torch.optim as optim

def get_finetune_optimizer(model):
    # 재학습 시에는 평소보다 10배 낮은 학습률로 시작하는 것이 관례
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
    return optimizer, scheduler

Example 5: 구조적 가지치기(Structured Pruning) 후 레이어 재구성

def prune_channels(model, layer_name, channel_indices):
    # 실제 텐서 크기를 줄여 연산 속도를 확보하는 방법
    layer = dict(model.named_modules())[layer_name]
    new_channels = layer.out_channels - len(channel_indices)
    
    # 가중치 복사 로직 (간략화)
    # 실제 구현 시에는 다음 레이어의 in_channels도 함께 수정해야 함
    print(f"Reducing {layer_name} to {new_channels} channels.")

Example 6: 영구적 마스크 적용 (Pruning Permanence)

def make_pruning_permanent(model):
    # 재학습 완료 후 마스크를 가중치에 완전히 적용하여 모델 크기 축소 준비
    for module in model.modules():
        if hasattr(module, 'weight_mask'):
            prune.remove(module, 'weight')
    print("Pruning masks are now permanent.")

Example 7: 희소 모델 평가 및 0이 아닌 파라미터 확인

def check_sparsity(model):
    total_zeros = 0
    total_elements = 0
    for name, buffer in model.named_buffers():
        if 'weight_mask' in name:
            total_zeros += torch.sum(buffer == 0).item()
            total_elements += buffer.nelement()
    
    sparsity = 100. * total_zeros / total_elements
    print(f"Global Sparsity: {sparsity:.2f}%")

5. 성능 회복을 위한 팁과 트릭

• Gradual Pruning: 한 번에 50%를 지우는 것보다 10%씩 5번에 걸쳐 지우고 중간마다 짧은 재학습을 거치는 것이 훨씬 안정적입니다.
• Batch Normalization 업데이트: 가지치기 후에는 데이터의 통계량이 변하므로 BN(Batch Norm) 레이어의 running mean/var를 다시 계산해야 합니다.
• L1 vs L2 Norm: 가중치 크기 기반 가지치기를 할 때, 0 근처의 값을 효과적으로 만드는 L1 정규화가 일반적으로 더 유리합니다.

6. 내용 출처

1. Frankle, J., & Carbin, M. (2018). The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks. ICLR.
2. Han, S., Pool, J., Tran, J., & Dally, W. (2015). Learning both Weights and Connections for Efficient Neural Network. NeurIPS.
3. PyTorch Official Documentation (torch.nn.utils.prune).