기술 요약: KLAS - 유사성을 활용한 신경망 스티칭을 통한 정확도-효율성 트레이드오프 개선

1. 문제 제기

대규모 사전 학습된 비전 트랜스포머 (ViT) 와 합성곱 신경망 (CNN) 의 배포는 근본적인 과제에 직면해 있습니다. 강력한 클라우드 서버부터 자원이 제한된 엣지 장치에 이르기까지 다양한 환경은 특정 컴퓨팅 예산 내에서 유연한 모델 선택을 요구합니다. 신경망 아키텍처 탐색 (NAS) 은 정확도와 효율성의 균형을 맞추려 시도하지만, 전통적인 1 대 1 NAS 는 계산 비용이 과도하게 높으며, 1 대 다 또는 제로샷 (zero-shot) NAS 접근법은 종종 단일 모델 설계 공간으로 제한되어 기존 사전 학습 모델들의 풍부한 풀을 활용하지 못합니다.

모델 스티칭 (stitching) 은 서로 다른 사전 학습된 '앵커 (anchor)' 모델의 블록들을 결합하여 보간된 네트워크를 생성함으로써 유망한 대안을 제시합니다. 그러나 SN-Net 과 같은 기존 다대다 (many-to-many) 스티칭 접근법은 인접한 앵커의 '가장 가까운 스티칭 (nearest stitching)'이나 고정된 블록 페어링과 같은 단순한 휴리스틱에 의존하여 스티칭 구성을 선택합니다. 이러한 휴리스틱은 종종 최적이지 않은 정확도 - 효율성 트레이드오프를 초래하며 다양한 모델 계열 간에 일반화되지 못합니다. 더 나아가, 스티칭 구성에 대한 탐색 공간은 방대합니다 ($k$개의 깊이 $n$인 모델에 대해 $O(k^2n^2)$). 따라서 포괄적인 탐색은 불가능합니다. 핵심 문제는 높은 학습 비용 없이 성능을 극대화할 수 있는 호환 가능한 스티칭 구성을 식별하는 원칙적이고 자동화된 방법의 부재입니다.

2. 방법론: KLAS 프레임워크

저자들은 사전 학습된 모델의 중간 표현 간의 유사성을 활용하여 스티칭 선택을 자동화하는 KLAS(KL 발산 기반 앵커 스티칭) 프레임워크를 제안합니다.

핵심 개념: 유사성 지표로서의 KL 발산

KLAS 는 성공적인 스티칭이 표현적 유사성(활성화의 기하학적 정렬) 과 기능적 유사성(작업 관련 정보 및 결정 경계의 보존) 모두에 의존한다고 가정합니다. 저자들은 클라이브 - 리블 (Kullback-Leibler, KL) 발산이 이 두 가지를 모두 포착하는 데 독특하게 적합하다고 주장합니다.

• 표현적 정렬: 활성화 분포 간의 통계적 거리를 정량화하여 경량 변환이 소스 활성화를 대상 활성화로 매핑할 수 있는지 여부를 나타냅니다.
• 기능적 정렬: 소스 블록의 출력이 대상 모델의 결정 경계를 얼마나 잘 보존하는지 평가하여, 스티칭된 네트워크가 최소한의 미세 조정으로 하류 작업 성능을 유지하도록 보장합니다.

스티칭 성공을 예측하는 데 종종 실패하는 Centered Kernel Alignment (CKA), 평균 제곱 오차 (MSE), 또는 교차 엔트로피 (CE) 와 같은 지표들과 달리, KL 발산은 선형 프로브 (linear probes) 에 의해 생성된 확률 분포의 호환성을 직접 측정합니다.

KLAS 알고리즘

이 프레임워크는 두 가지 주요 단계로 작동합니다.

1. 앵커 선택:

   • KLAS 는 후보 사전 학습 모델의 마지막 블록의 소프트맥스 확률 분포 간의 KL 발산을 계산합니다.
   • 마지막 블록 KL 발산이 낮은 모델은 유사한 결정 경계와 클래스 신뢰도 분포를 공유하므로 호환 가능한 앵커로 간주됩니다. 이는 이전 연구에서 사용된 휴리스틱인 '가장 가까운 스티칭'을 대체합니다.

2. 블록 선택:

   • 앵커가 선택되면 KLAS 는 스티칭 점수 $\Gamma(i, j)$를 사용하여 블록 쌍을 평가합니다. 이는 다음과 같이 정의됩니다:
     $$\Gamma(i, j) = \frac{\Theta(P^f_i, P^g_j)}{\Theta(P^g_j, P^g_{j+1})}$$
     여기서:
     • $\Theta(P^f_i, P^g_j)$는 교차 앵커 활성화 거리 (소스 블록 $i$와 대상 블록 $j$ 간의 KL 발산) 입니다.
     • $\Theta(P^g_j, P^g_{j+1})$는 앵커 내 블록 용량 (대상 앵커 내 연속된 블록 간의 KL 발산) 입니다.
   • 이 점수는 가상의 스티칭 레이어에 가해지는 변환 부담과 대상 블록이 분포 차이를 흡수할 수 있는 능력 사이의 균형을 맞춥니다.
   • ProbeNet: 전체 스티칭 레이어를 학습하지 않고 이러한 점수를 효율적으로 계산하기 위해 KLAS 는 ProbeNet 이라는 통합 아키텍처를 사용합니다. 이는 앵커 모델의 모든 블록에 선형 분류기 프로브를 연결하고 단일 순전파 - 역전파 과정에서 함께 학습합니다. 이는 개별 프로브를 학습하는 것에 비해 무시할 수 있는 일회성 비용 (예: Swin-B 의 경우 0.25 GPU 일) 을 초래합니다.

3. 구성 선택:

   • 후보 스티칭 구성은 FLOPs 버킷 내 최소 점수에 대한 조정 가능한 임계값 $\tau$에 따라 필터링됩니다. 이는 정확도 - 효율성 스펙트럼 전반에 걸쳐 배포 가능한 모델들의 밀집된 파레토 프론트 (dense Pareto front) 를 보장합니다.

3. 주요 기여

1. 휴리스틱의 한계 식별: 이 논문은 SN-Net 과 같은 단순한 휴리스틱이 보편적으로 최적이지 않으며, 휴리스틱이 간과하는 우수한 '멀리 떨어진 스티칭 (far stitching)' 구성 (예: Swin-Ti 를 직접 Swin-B 에 연결) 을 종종 놓친다는 것을 보여줍니다.
2. 원칙적인 유사성 지표: 스티칭 레이어의 인스턴스화나 학습 없이 표현적 및 기능적 정렬을 모두 포착하는 스티칭 선택을 위한 지표로서 KL 발산을 도입했습니다.
3. KLAS 프레임워크: 다양한 모델 계열 (ViT, CNN, LLM) 전반에 걸쳐 앵커 및 블록 선택을 자동화하는 일반화 가능하고 저비용의 프레임워크입니다.
4. 실증적 검증: 여러 데이터셋 (ImageNet-1K, CIFAR-100, ADE20K, TruthfulQA) 과 아키텍처에 걸쳐 베이스라인 대비 일관된 개선을 보여주는 포괄적인 실험 결과입니다.

4. 실험 결과

저자들은 동일한 미세 조정 조건 (비전의 경우 50 에포크, LLM 의 경우 10 에포크) 하에서 KLAS 를 SN-Net 및 기타 베이스라인과 비교 평가했습니다.

• 정확도 - 효율성 트레이드오프:

  • ImageNet-1K: KLAS 는 SN-Net 대비 동일한 계산 비용에서 최대 1.21% 높은 Top-1 정확도를 달성하거나, 동등한 정확도에서 FLOPs 를 1.33 배 감소시켰습니다.
  • AUC 점수: Swin 계열에서 KLAS 는 0.8950의 곡선 아래 면적 (AUC) 을 달성하여 SN-Net (0.8345) 및 CKA (0.8124) 와 같은 기타 유사성 지표를 크게 능가했습니다.
  • 교차 아키텍처: KLAS 는 ResNet-18 과 Swin-B 를 성공적으로 스티칭하여 유사한 FLOPs 에서 캐스케이드 추론 베이스라인보다 2-3% 높은 Top-1 정확도를 달성했습니다.

• 밀집 예측 (시맨틱 세그멘테이션):

  • Mask2Former 를 사용한 ADE20K 에서 KLAS 가 선택한 스티칭은 유사한 FLOPs 예산에서 SN-Net 베이스라인 대비 최대 +0.9% mIoU를 달성했습니다.

• 대형 언어 모델 (LLM):

  • TruthfulQA 벤치마크에서 Llama 1B 및 3B 앵커를 사용할 때, KLAS 스티칭 모델은 일부 구성에서 더 적은 매개변수를 가졌음에도 휴리스틱 기반 LLM 스티칭 방법인 ESTA 를 최대 0.017 ROUGE-1 및 0.033 ROUGE-2만큼 능가했습니다.

• 지표 효과성:

  • KL 발산은 가능한 Swin Ti-B 스티칭 구성의 **88.9%**를 성공적으로 복원한 반면, CKA 는 5.5%, SN-Net 휴리스틱은 61.1% 만 복원했습니다.

5. 중요성 및 주장

이 논문은 KLAS 가 임시방편적인 휴리스틱을 원칙적이고 유사성 기반의 접근법으로 대체함으로써 모델 스티칭의 중요한 격차를 해소한다고 주장합니다. 그 중요성은 다음과 같습니다.

• 일반화 가능성: 단일 계열로 제한된 방법의 한계를 극복하고 다양한 아키텍처 계열 (트랜스포머, CNN, 하이브리드) 과 모달리티 (비전, 언어) 전반에 걸쳐 작동합니다.
• 비용 효율성: 전통적인 NAS 의 prohibitive 한 비용을 피하면서 프로브 학습과 스티칭 레이어 미세 조정만으로도 배포 가능한 모델들의 연속적인 파레토 프론트를 생성하며, 추가적인 계산 오버헤드는 무시할 수준입니다.
• 체계적 최적화: KL 발산을 통해 분포 차이를 명시적으로 포착함으로써 휴리스틱 방법이 놓치는 직관적이지 않고 고성능을 발휘하는 스티칭 구성 (예: '멀리 떨어진 스티칭') 을 발견할 수 있게 합니다.

저자들은 KLAS 가 처음부터 학습할 필요 없이 기존 사전 학습 모델 생태계를 효과적으로 활용하여 배포 시나리오에서 정확도 - 효율성 트레이드오프를 최적화하는 견고하고 자동화된 솔루션을 제공한다고 결론지었습니다.