지난번에 INT8 양자화로 크기를 92% 줄여봤다. INT8로 바꾸니까 크기 75% 줄고, 속도 4배 빨라지고, 정확도는 그대로였다. 너무 잘 돼서 의심스러울 정도였다.

그래서 다른 경량화 기법도 해보기로 했다. 프루닝이랑 지식 증류. 둘 다 이름은 많이 들어봤는데 직접 해본 적은 없었다.

결론부터 말하면, 양자화만큼 쉽지 않았다. 프루닝은 50%를 날리고도 오히려 느려졌고, 지식 증류는 정확도 13%p를 내줬다. 두 기법이 왜 교과서대로 안 가는지, 그리고 그럼 언제 써야 하는지 정리했다.

프루닝: 불필요한 가중치 잘라내기

프루닝은 별거 없다. 딥러닝 모델의 가중치 중에서 값이 작은 것들은 결과에 별로 영향을 안 준다. 그러니까 그냥 0으로 만들어버리자. Pruning이라는 단어 자체가 "가지치기"라는 뜻이다. 식물 가지치기에서 그대로 가져온 용어다.

크게 두 가지 방식이 있다.

Unstructured Pruning은 개별 weight를 하나씩 0으로 만든다. 유연하긴 한데, 실제로 속도가 빨라지려면 하드웨어가 sparse 연산을 지원해야 한다.

Structured Pruning은 뉴런이나 채널 단위로 통째로 잘라낸다. 네트워크 구조 자체가 바뀌니까 확실히 빨라지는데, 정확도 손실이 더 크다.

일단 Unstructured로 해보기로 했다. TensorFlow Model Optimization 라이브러리에서 지원하니까.

프루닝 알고리즘 선택

프루닝 알고리즘도 여러 가지가 있다.

One-shot Pruning은 한 번에 확 잘라버린다. 구현은 가장 간단한데 정확도가 많이 깎인다.

Iterative Magnitude Pruning은 조금 자르고 Fine-tuning, 또 조금 자르고 Fine-tuning을 반복한다. Lottery Ticket 논문(Frankle & Carbin, 2019)에서 쓴 방식이다. 정확도는 좋은데 학습 시간이 10배쯤 뛴다.

Automated Gradual Pruning(AGP)은 학습하면서 점진적으로 자른다. 3차 함수 스케줄로 초반엔 많이, 후반엔 조금씩. 하이퍼파라미터가 적어서 TensorFlow Model Optimization이 기본값으로 권장한다.

시간 대비 효과만 보고 AGP를 골랐다.

50% 가중치 제거 실험

목표는 가중치의 50%를 0으로 만드는 거다. 절반을 날려도 정확도가 유지되면 대성공이다.

pruning_params = {
    'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
        initial_sparsity=0.0,   # 0%에서 시작
        final_sparsity=0.5,     # 50%까지 점진적 증가
        begin_step=0,
        end_step=3000           # 3000 스텝에 걸쳐 증가
    )
}

pruned_model = tfmot.sparsity.keras.prune_low_magnitude(
    base_model, 
    **pruning_params
)

base 모델을 10 에포크 학습시켜서 76.03% 정확도를 만들고, 거기에 프루닝을 적용하면서 10 에포크 더 Fine-tuning 했다.

결과가 좀 의외였다.

50% 잘라냈는데 정확도가 3%p 올라간다. 잘못 측정한 줄 알았는데 반복해도 같았다.

중요도 낮은 weight를 치우면 regularization처럼 작용해서 overfitting이 줄어든다. 거기에 Fine-tuning 10 에포크가 얹혀서 남은 weight가 더 조여진 영향도 있을 거다. CIFAR-10이 비교적 단순한 데이터셋이라서 이게 크게 먹혔다. 복잡한 데이터에서는 다르게 나올 수 있다.

그런데 파일 크기가 안 줄었다

H5 파일로 저장해봤더니 크기가 똑같았다. 13MB 그대로.

왜 그런지 찾아보니까, H5 형식은 sparse matrix를 지원하지 않는다. 0인 weight도 그냥 0.0이라는 값으로 저장한다. 공간을 똑같이 차지하는 거다.

H5 형식:
[1.2, 0.0, 0.0, 3.4, 0.0, 0.0, 2.1, 0.0, ...]
→ 모든 값이 4바이트씩 저장됨

TFLite로 변환하면 압축이 된다고 해서 해봤다. 4.3MB. FP32 TFLite와 똑같다. 양자화까지 적용해야 줄어드는 것 같다.

INT8 양자화를 추가로 적용했더니 1.1MB가 됐다. 근데 이건 프루닝 안 한 모델도 마찬가지다. 프루닝 자체로는 크기 감소 효과가 없는 셈이다.

속도까지 반대로 갔다

크기는 그렇다 치고, 속도라도 빨라졌을 거라고 기대했다. 50%를 계산 안 해도 되니까.

측정해봤다.

약 1.4배 느려졌다. 드라마틱한 5배는 아니지만, 계산량을 반으로 줄였는데 오히려 느려진 건 맞다. 이론대로면 2배는 빨라져야 한다. 한 줄로 정리하면 FLOPs는 절반으로 줄었는데 실제 latency는 오히려 늘었다.

한참을 고민했다. 왜 이런 결과가 나왔을까.

결국 문제는 런타임이다. Dense 연산은 메모리를 쭉 연속으로 읽으니까 캐시가 잘 붙는다. sparse 연산은 0이 아닌 값만 골라야 하니까 조건문이 끼고, 메모리 접근도 들쭉날쭉이다.

# Dense
for i in range(N):
    result += weight[i] * input[i]

# Sparse (개념적으로)
for i in range(N):
    if weight[i] != 0:
        result += weight[i] * input[i]

조건 체크 오버헤드와 캐시 미스가 같이 온다. 프루닝의 문제가 아니라 런타임의 문제다. sparse 연산을 네이티브로 처리하는 하드웨어/라이브러리가 없으면 오히려 느려진다.

추가로 짚을 게 하나 있다. 프루닝 후에는 tfmot.sparsity.keras.strip_pruning()을 호출해서 학습용 wrapper를 제거해야 한다. 이걸 빼먹으면 추론 시에도 wrapper 오버헤드가 남는다. 이 영향도 느려지는 데 한몫했을 수 있다.

한 줄로 정리하자면, 50% 계산량을 날리고도 역효과가 났다.

정확도는 78.93%로 올랐지만, 이건 보너스에 가깝다. 목표였던 크기는 그대로였고, 속도는 오히려 1.4배 느려졌으니까. 논문에서 말하는 속도 이점을 실제로 얻으려면 sparse 연산을 네이티브로 지원하는 하드웨어/런타임이 있어야 한다. 일반 CPU + TFLite 조합에서는 프루닝은 "썼다"와 "효과 봤다" 사이의 간격이 크다.

지식 증류: 선생님 모델의 지식을 전달하기

기분을 바꿔서 지식 증류를 해보기로 했다.

크고 정확한 Teacher 모델을 먼저 학습시킨다. 그 다음에 작은 Student 모델이 Teacher를 따라하면서 학습한다. 혼자 학습하는 것보다 Teacher의 "지식"을 전달받아서 더 좋은 성능을 낼 수 있다.

근데 Student가 대체 뭘 배우는 걸까.

Hard Target vs Soft Target

일반 학습은 정답 레이블을 쓴다. 고양이면 [0, 0, 1, 0, 0, 0, 0, 0, 0, 0]. Hard Target이다.

지식 증류는 Teacher의 출력을 쓴다. [0.02, 0.01, 0.70, 0.15, 0.05, ...]. "고양이 70%, 개 15%, 사슴 5%". 이게 Soft Target이다.

Hard Target이 "고양이다" 한 마디라면, Soft Target은 "고양이인데 개랑 좀 비슷하고 자동차랑은 전혀 다르다"까지 담고 있다. Teacher가 학습하면서 쌓은 클래스 간 거리 감각이 이 확률 분포에 녹아 있다. Student는 그걸 따라가면서 Teacher의 판단 방식을 배운다.

Temperature: 확률 분포를 부드럽게

근데 문제가 있다. Teacher가 너무 확신하면 Soft Target도 Hard Target이랑 비슷해진다.

[0.99, 0.01, 0.00, 0.00, ...]

99%면 사실상 "고양이다!"랑 다를 게 없다. 다른 클래스 정보가 거의 없다.

그래서 Temperature라는 걸 쓴다. Softmax 연산에서 로짓을 T로 나눠준다. T가 크면 확률 분포가 평평해진다.

T=1 (일반):  [0.66, 0.24, 0.10]
T=3 (높음):  [0.44, 0.32, 0.24]
T=5 (매우):  [0.40, 0.33, 0.27]

학습할 때는 T=3 정도로 부드럽게 만들어서 정보를 최대한 전달하고, 실제 추론할 때는 T=1로 명확하게 예측한다.

두 가지 Loss 조합

지식 증류의 Loss는 두 가지를 섞는다.

한쪽은 Student 예측과 정답 레이블 차이(Hard Loss, Cross-Entropy). "정답 방향"을 알려준다. 다른 한쪽은 Student와 Teacher Soft Target의 차이(Soft Loss, KL Divergence). "Teacher 따라하기"다. 두 Loss를 가중치로 묶어서 같이 학습시킨다.

Total Loss = α × Hard Loss + (1-α) × Soft Loss

보통 α=0.1을 쓴다. Hard Loss 10%, Soft Loss 90%. Teacher의 지식을 주로 배우되, 정답도 참고하는 식이다.

Hinton 논문은 "Hard Loss 쪽 가중치를 훨씬 작게 주면 좋은 결과가 나왔다"고만 언급한다. 실험에 따라 구체적 최적값은 다르니까, 표준 관행처럼 쓰이는 α=0.1을 초깃값으로 잡고 필요하면 조정하는 식으로 쓰면 된다.

Teacher와 Student 설계

Teacher는 크고 복잡하게 만들었다. Conv 레이어 6개에 BatchNorm 넣고, Dense도 512 유닛. 파라미터 약 320만 개. CIFAR-10 기준 80% 이상 나올 수 있는 최소 복잡도를 노린 거다.

한 가지 설계 포인트가 있다. 두 모델 모두 마지막 Dense 레이어에 softmax를 빼고 logits로 출력하게 했다. 이유는 뒤에 증류 학습 구현에서 같이 설명한다.

def create_teacher_model():
    """Teacher 모델 - ~3.2M 파라미터"""
    model = keras.Sequential([
        keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
        keras.layers.BatchNormalization(),
        keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
        keras.layers.MaxPooling2D(),
        # ... Block 2, 3 ...
        keras.layers.Flatten(),
        keras.layers.Dense(512, activation='relu'),
        keras.layers.Dense(10, activation=None)  # logits 출력
    ])
    return model

Student는 반대로 모바일에 올릴 수 있는 선에서 최대한 줄였다. Conv 레이어 2개에 Dense 128 유닛. 파라미터 약 27만 개. Teacher의 1/12 크기다.

def create_student_model():
    """Student 모델 - ~270K 파라미터"""
    model = keras.Sequential([
        keras.layers.Conv2D(16, 3, padding='same', activation='relu'),
        keras.layers.MaxPooling2D(),
        keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
        keras.layers.MaxPooling2D(),
        keras.layers.Flatten(),
        keras.layers.Dense(128, activation='relu'),
        keras.layers.Dense(10, activation=None)  # logits 출력
    ])
    return model

증류 학습 구현

TensorFlow에는 증류용 API가 따로 없다. train_step을 직접 짜야 한다.

여기서 앞에서 남긴 숙제를 풀고 가자. 두 모델의 마지막 Dense를 softmax 없이 logits로 출력하게 한 이유다. 증류 loss는 softmax(logits / T)를 넣어야 Temperature가 제 역할을 한다. 만약 모델이 이미 softmax가 적용된 확률을 반환하면, 그 위에 다시 softmax(확률 / T)를 씌우는 꼴이 된다. 확률은 이미 0~1 범위고 그걸 T로 나눠봤자 범위가 더 좁아질 뿐이라, 한 번 더 softmax를 태우면 분포가 거의 균등(uniform)에 가까워진다. Soft Target의 클래스 간 정보가 사라지는 거다.

그래서 모델은 logits를 뱉고, 학습 loss 쪽만 from_logits=True로 받는다.

class DistillationModel(keras.Model):
    def train_step(self, data):
        x, y = data

        # Teacher 예측 (학습 안함) — logits 반환
        teacher_predictions = self.teacher(x, training=False)

        with tf.GradientTape() as tape:
            # Student 예측 — logits 반환
            student_predictions = self.student(x, training=True)

            # Hard Loss (student_loss_fn은 from_logits=True로 생성)
            student_loss = self.student_loss_fn(y, student_predictions)

            # Soft Loss (logits / T → softmax)
            distillation_loss = self.distillation_loss_fn(
                tf.nn.softmax(teacher_predictions / self.temperature),
                tf.nn.softmax(student_predictions / self.temperature)
            ) * (self.temperature ** 2)

            # 최종 Loss
            total_loss = self.alpha * student_loss + (1 - self.alpha) * distillation_loss

        # Student만 업데이트
        gradients = tape.gradient(total_loss, self.student.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.student.trainable_variables))

        return {"total_loss": total_loss, ...}

T²를 곱하는 이유는 기울기 크기를 보정하기 위해서다. Temperature로 나누면 기울기가 작아지니까 T²로 보상해준다.

실험 결과

Teacher를 20 에포크 학습시켜서 82.79% 정확도를 만들었다.

Student를 두 가지 방식으로 학습시켰다. 하나는 단독으로 20 에포크 학습. 다른 하나는 Teacher한테 증류 받으면서 20 에포크 학습.

Teacher:         82.79%
Student (단독):  67.08%
Student (증류):  69.76%

증류 효과는 +2.68%p. 단독 학습보다 확실히 나았다.

근데 Teacher랑 비교하면 13%p나 낮다. 파라미터 1/12로 줄였으니 담을 수 있는 정보량 자체가 빠진다. 그런데 Transformer 쪽은 DistilBERT가 레이어를 반 토막 내도 원본 97% 성능을 유지한다(뒤에서 다시 얘기). 증류가 모델 구조에 따라 결과가 크게 갈린다는 얘기다.

지식 증류의 trade-off

파라미터는 확실히 줄었다. Teacher가 약 325만 개(H5 기준 38MB), Student가 약 27만 개(H5 기준 1.05MB). 파라미터 기준 91.7% 감소다. 실제 배포용 INT8 TFLite로 바꾸면 Student가 529KB까지 내려간다.

속도도 빨라졌다. Teacher 23.18ms, Student 19.95ms. 다만 파라미터가 1/12인데 속도는 그만큼 줄진 않았다. 레이어 수가 적어서 병렬 처리 이득이 제한적인 것 같고, 배치 크기나 모바일 환경에서는 또 다를 수 있다.

정확도 손실이 문제다. 13%p면 적지 않다. 응용 분야에 따라서는 치명적일 수도 있다.

그리고 학습 시간이 2배다. Teacher 20 에포크 + Student 20 에포크. 단독으로 Student만 40 에포크 학습시키면 어떨까 싶기도 하다. 그것도 해봐야 할 것 같다.

세 가지 경량화 기법 비교

CIFAR-10 실험 결과를 정리하면 이렇다.

모델 정확도 파라미터 TFLite 크기 TFLite 속도
Teacher (원본) 82.79% 3.25M — —
Base (중간) 77.88% 1.12M 4,367 KB 754.7μs
+ FP16 양자화 77.88% 1.12M 2,188 KB 755.4μs
+ INT8 양자화 77.88% 1.12M 1,108 KB 169.8μs (4.44x)
+ 프루닝 50% 78.93% 1.12M 1,104 KB 238.4μs
Student (단독) 67.08% 0.27M — —
Student (증류) 69.76% 0.27M 529 KB 68.0μs

INT8 양자화는 가성비 최고다. 5분이면 구현이 끝나고 효과가 확실하다.

프루닝은 속도가 느려진 게 문제다. 드라마틱한 저하는 아니지만, 계산량을 반으로 줄이고 오히려 느려졌다는 자체가 의미가 없다는 뜻이다. 하드웨어 지원 없이는 실용성이 떨어진다.

지식 증류는 극단적인 경량화가 필요할 때 쓸 수 있지만 정확도 타협이 크다.

모델마다 주류가 다르다

이건 전제를 하나 깔고 봐야 한다. 내가 실험한 건 CNN이다. Transformer나 LLM에 그대로 적용되는 결론이 아니다.

CNN에서는 양자화가 제일 깔끔하게 먹힌다. Conv 연산이 단순해서 INT8 최적화가 잘 붙고, 필터 간 중복이 많아서 프루닝 여지도 있다(하드웨어가 받쳐주면).

Transformer 쪽은 무게 중심이 다르다. DistilBERT가 대표 사례인데, BERT의 12개 레이어를 6개로 줄이면서 GLUE 벤치마크에서 원본 대비 97% 성능을 유지했다. 내 CNN 실험에서 1/12로 줄였을 때 13%p가 빠진 것과 비교하면 격차가 크다. 레이어 단위로 반 토막을 내도 버티는 구조 자체가 증류와 궁합이 좋다는 뜻이다.

그래서 요즘 LLM 경량화 흐름도 증류가 중심이다. GPT-4 급 모델의 지식을 작은 모델로 옮기고, 거기에 INT8 양자화를 한 번 더 얹는다. 경량화 기법은 서로 배타적인 게 아니라 쌓아 올리는 거다. 모델 종류마다 어느 층을 먼저 쌓는지가 달라질 뿐이다.

결론

셋을 다 해보니 선택의 기준이 분명해졌다.

양자화는 거의 무조건이다. 구현 쉽고, 효과 확실하고, 손실도 미미했다. 안 할 이유가 없다. 반면 프루닝은 이번 실험에서 본 것처럼 하드웨어와 런타임이 sparse 연산을 네이티브로 안 받치면 역효과가 난다. 조건이 맞는지 확인 못 한 상태로 꺼내면 계산량만 줄어들고 속도는 오히려 깎인다.

지식 증류는 다른 축이다. 크기를 극단적으로 줄여야 할 때, 그리고 모델 구조를 바꿀 수 있을 때 의미가 있다. CNN에서는 13%p 손실을 감수해야 했지만 Transformer 계열이면 얘기가 달라진다. 정확도 손실을 어디까지 감수할 수 있는지부터 정해야 한다.

내가 쓰는 순서는 단순하다. 양자화 먼저 — 거의 공짜라서. 그래도 부족하면 증류. 프루닝은 마지막 선택지이자, 조건 맞을 때만 꺼낸다. 보이스피싱 탐지 KoBERT라면 양자화 + DistilKoBERT 조합이 현실적인 첫 시도가 될 것 같다.

경량화 기법은 정리됐다. 다음 글에서는 INT8 모델을 Android에 올려서 진짜로 돌아가는지 확인해본다

소스 코드

https://github.com/yjcho9317/CIFAR10_OnDevice

참고 자료

  • Hinton et al. (2015) - Distilling the Knowledge in a Neural Network
  • Frankle & Carlin (2019) - Lottery Ticket Hypothesis
  • Zhu & Gupta (2017) - To prune, or not to prune
  • Sanh et al. (2019) - DistilBERT
  • HMG Developers 프루닝 블로그: https://developers.hyundaimotorgroup.com/blog

'AI' 카테고리의 다른 글

온디바이스 AI 경량화 (3) — LiteRT로 INT8 모델 Android 배포와 GPU Delegate의 함정  (0) 2026.04.17
온디바이스 AI 경량화 (1) — INT8 양자화로 CIFAR-10 모델 92% 줄이기 (TFLite PTQ)  (0) 2026.04.17
MLOps 구축기 (5) - MLflow 기반 MLOps 파이프라인 전체 정리  (1) 2026.03.13
MLOps 구축기 (4) - GitHub Actions CI/CD와 Prometheus + Grafana 모니터링  (0) 2026.03.13
MLOps 구축기 (3) - MLflow 모델을 FastAPI로 서빙하기  (0) 2026.03.08

티스토리툴바