학습 가능한 채널 순열을 통한 트랜스포머 구조적 희소성 최적화

초록

본 논문은 트랜스포머 모델에 적용되는 N:M 구조적 희소성(pruning) 성능을 높이기 위해, 가중치 채널을 학습 가능한 방식으로 순열하는 프레임워크를 제안한다. 비용 행렬을 학습하고, 이를 기반으로 차분 가능한 이분 매칭(Sinkhorn) 솔버를 이용해 최적 이진 순열 행렬을 도출하며, 순열과 희소성 마스크를 공동으로 최적화하는 손실 함수를 설계한다. 비전·언어 트랜스포머에 적용한 실험에서 기존 휴리스틱 기반 순열 방법을 크게 능가하는 결과를 보인다.

상세 분석

이 논문은 구조적 희소성, 특히 N:M(예: 2:4) 제약 하에서 트랜스포머 모델의 정확도 저하를 최소화하기 위한 새로운 접근법을 제시한다. 핵심 아이디어는 “채널 순열”이다. 기존 연구에서는 중요도 점수를 기반으로 그리디하게 채널을 재배열했지만, 이는 전역적인 최적화를 보장하지 못하고 계산 비용이 크게 증가한다. 저자는 이를 해결하기 위해 세 가지 주요 구성 요소를 설계하였다. 첫째, 학습 가능한 순열 비용 행렬(Cost Matrix) 은 각 입력 채널 i를 출력 위치 j에 매핑했을 때 발생하는 비용을 실시간으로 추정한다. 이 비용은 가중치의 saliency, N:M 마스크와의 정렬 정도, 그리고 지식 증류(distillation) 손실 등을 종합한 형태로, 파라미터화된 신경망(예: 작은 MLP)으로 구현된다. 둘째, 미분 가능한 이분 매칭 솔버는 비용 행렬을 입력으로 받아 최적의 이진 순열 행렬 P를 생성한다. 여기서는 Birkhoff 다면체 위에서 엔트로피 정규화된 최적화를 수행하고, Sinkhorn 반복을 통해 근사적인 doubly‑stochastic 행렬을 얻은 뒤, Straight‑Through Estimator(STE)를 이용해 이진화한다. 이렇게 하면 순열 연산이 그래디언트 흐름을 차단하지 않으며, 전체 파이프라인을 end‑to‑end로 학습할 수 있다. 셋째, 희소성 최적화 손실은 (1) 기본 작업 손실(예: 교차 엔트로피), (2) 지식 증류 손실(teacher‑student 출력 정합), (3) 순열 비용 자체를 최소화하는 정규화 항으로 구성된다. 이 손실은 순열과 마스크가 동시에 최적화되도록 유도한다.
구조적 관점에서, 트랜스포머의 MHA와 FFN 블록은 여러 선형 변환이 병렬로 연결돼 있어 채널 순열이 서로 얽힌다. 논문은 “바인딩” 규칙을 제시해, 현재 레이어의 입력 채널 순열이 이전 레이어의 출력 채널 순열과 일관되게 적용되도록 설계하였다. 이는 활성값을 별도로 재배열할 필요 없이 가중치만 재배열함으로써 런타임 오버헤드를 최소화한다. 또한, 대규모 모델에 대한 메모리·연산 부담을 줄이기 위해 그룹‑와이즈 순열을 도입, 채널을 고정 크기 G로 나누어 각 그룹 내에서만 순열을 학습한다. 이는 비용 행렬의 차원을 크게 감소시켜 Sinkhorn 연산을 효율화한다.
실험에서는 ViT‑B/16, LLaMA‑7B, 그리고 멀티모달 VLM 등 다양한 아키텍처에 적용했으며, 2:4 및 4:8 희소성 비율에서 기존 그리디 순열(Pool & Yu, Plug‑and‑Play) 대비 Top‑1 정확도 향상을 0.51.2%p 달성했다. 특히, 큰 언어 모델에서는 파라미터 감소율이 동일함에도 불구하고 퍼플렉시티 상승이 현저히 낮았다. Ablation study는 (a) 비용 행렬을 학습하지 않고 무작위 초기화할 경우 성능이 급격히 떨어짐을, (b) Sinkhorn 반복 수를 510회로 제한해도 충분히 좋은 근사치를 얻을 수 있음을, (c) 그룹‑와이즈 크기 G가 64~128일 때 최적의 트레이드오프를 보임을 보여준다.
한계점으로는 (1) 비용 행렬 학습이 추가 파라미터와 메모리를 요구한다는 점, (2) 현재는 N:M 마스크와 결합된 정적 프루닝에만 적용했으며, 동적 마스크(예: 스파스 레이어 재학습)와의 연계는 미탐색이다. 향후 연구에서는 비용 행렬에 메타‑학습을 도입하거나, 순열을 사전‑학습된 프리트레인 단계에 통합해 전체 파이프라인을 더욱 가볍게 만드는 방향이 기대된다.