기술 요약: MoECodec

문제 정의

기계 지향적 이미지 압축(Machine-oriented image compression)은 분류, 탐지, 세그멘테이션과 같은 다운스트림 컴퓨터 비전 작업을 지원하기 위한 핵심 연구 분야로 부상했습니다. 그러나 기존 방식들은 다음과 같은 중대한 한계에 직면해 있습니다:

1. 태스크 특화 설계 (Task-Specific Designs): 특정 태스크를 위한 엔드 투 엔드(end-to-end) 최적화는 전용 인코더-디코더 쌍이나 멀티 브랜치 아키텍처를 요구하며, 이는 상당한 파라미터 오버헤드와 배포 복잡성을 초래합니다.
2. 전이 기반 적응 (Transfer-Based Adaptation): 소수의 학습 가능한 파라미터(예: 어댑터)를 사용하여 인간 지향적 베이스 코덱을 적응시키는 방법들은 모델 내부에 본질적으로 학습되기보다는 외부적으로 부착되어 있으며, 태스크 특화된 휴리스틱에 의존합니다.
3. 정적 계산 패턴 (Static Computation Patterns): 두 가지 흐름 모두 대부분의 이미지 토큰에 대해 균일한 변환을 적용합니다. 이는 서로 다른 이미지 영역이 기계 인지 측면에서 현저히 다른 의미적 중요성과 복잡성을 가진다는 사실을 고려하지 못하므로, 자원 할당의 비효율성을 초래합니다.

방법론

저자들은 트랜스포머 기반의 학습된 이미지 압축(Learned Image Compression, LIC) 모델에 토큰 단위의 Mixture-of-Experts(MoE)를 통합하는 통합 이미지 압축 프레임워크인 MoECodec을 제안합니다. 핵심 혁신은 트랜스포머 블록의 표준 피드 포워드 네트워크(FFN) 레이어를 MoE 레이어로 대체하여, 입력 콘텐츠와 태스크 목적에 따라 동적인 토큰 수준의 계산을 가능하게 하는 것입니다.

주요 아키텍처 구성 요소

1. 토큰 단위 MoE 통합 (Token-Wise MoE Integration):

   • MoECodec은 트랜스포머 기반 LIC 백본(예: TIC, TinyLIC, TCM)의 인코더와 디코더 모두에서 FFN 레이어를 대체합니다.
   • 입력 토큰은 전문화된 전문가 네트워크(expert networks)로 적응적으로 라우팅되어, 단일 통합 모델 내에서 이질적인 지각 목적을 지원할 수 있게 합니다.

2. 라우팅 전략 (Expert-Choice + Spatial Regularization):

   • 과제: 표준 트랜스포머 기반 LIC 모델은 종종 미세한 토큰화(예: 1×1 패치)를 사용합니다. 단순한 토큰 선택(Token-Choice, TC) 라우팅은 공간적 문맥의 부재로 인해 노이즈가 많고 파편화된 전문가 할당(솔트 앤 페퍼 아티팩트)을 초래합니다.
   • 해결책: 저자들은 각 전문가가 전체 시퀀스에서 자신이 선호하는 상위 $k$개의 토큰을 능동적으로 선택하는 Expert-Choice (EC) 라우팅을 채택합니다. 이는 결정의 관점을 역전시켜, 전문가가 글로벌 문맥을 고려하고 보조적인 로드 밸런싱 손실 없이도 균형 잡힌 활용을 보장하도록 합니다.
   • 공간 총 분산 (Spatial Total Variance, TV) 정규화: EC 라우팅의 국부적 공간 불일치를 추가로 해결하기 위해, 훈련 손실에 공간 TV 정규화 항을 추가합니다. 이는 공간 전문가 친화도 맵(affinity maps)의 고주파 변동에 패널티를 부여하여, 조각별로 매끄럽고 일관된 전문가 할당을 유도합니다.

3. 경량 전문가 아키텍처 (Group Shuffle MLP):

   • 과제: FFN을 $E$개의 독립적인 전문가로 교체하면 일반적으로 $E\times$의 파라미터 오버헤드가 발생하며, 이는 FFN이 파라미터 수를 지배하는 압축 모델에서 매우 부담스러운 수치입니다. 단순히 전문가의 크기를 줄이는 것은 채널 집계 능력을 저하시킵니다.
   • 해결책: 저자들은 **Group Shuffle MLP (GShMLP)**를 제안합니다. 이 아키텍처는 입력 채널을 $G$개의 그룹으로 분할하고 각 그룹 내에서 독립적인 선형 투영을 적용합니다. 두 투영 레이어 사이에 파라미터가 없는 채널 셔플(channel shuffle)을 삽입하여 교차 채널 상호작용을 복원합니다. 이 설계는 효과적인 채널 집계를 유지하면서 파라미터 복잡도를 크게 줄입니다.

주요 기여

1. 통합 프레임워크: MoECodec은 토큰 단위 MoE를 트랜스포머 기반 코덱에 통합한 최초의 통합 LIC 프레임워크로, 단일 모델이 재구성 품질과 다수의 기계 다운스트림 태스크를 공동으로 최적화할 수 있게 합니다.
2. LIC 특유의 과제 해결: 본 논문은 MoE를 LIC에 도입할 때 발생하는 두 가지 구체적인 과제를 식별하고 해결했습니다:
   • 공간적 불연속성: Expert-Choice 라우팅과 공간 TV 정규화를 통해 해결되었습니다.
   • 파라미터 오버헤드: 경량 GShMLP 아키텍처를 통해 해결되었습니다.
3. 동적 계산 할당: 이 프레임워크는 이질적인 태스크 목적이 본질적으로 수용되는 동적 계산 패러다임을 구축하여, 모델이 태스크 관련 영역에 계산을 집중할 수 있도록 합니다.

실험 결과

광범리한 실험을 통해 재구성과 기계 태스크 모두에서 강력한 베이스라인 대비 일관된 개선을 입증했습니다.

• 이미지 재구성: 세 가지 트랜스포머 기반 백본(TIC, TinyLIC, TCM)에 적용했을 때, MoECodec은 Kodak 데이터셋(PSNR 기준)에서 각각 11.54%, 8.12%, **6.42%**의 BD-Rate 개선을 달야하는 동시에 전체 파라미터 수를 줄였습니다.
• 기계 다운스트림 태스크: TIC를 베이스로 사용했을 때, MoECodec은 Full Finetune, TransTIC, AdaptICMH 베이스라인보다 우수한 성능을 보였습니다:
  • 분류 (ImageNet-1K): 7.74% BD-Rate 개선.
  • 객체 탐지 (COCO): 8.98% BD-Rate 개선.
  • 인스턴스 세그멘테이션 (COCO): 9.72% BD-Rate 개선.
• 효율성: MoECodec은 활성화된 파라미터를 7.51M에서 6.25M으로 줄였으며, RTX 4090에서 142ms의 추론 지연 시간을 달성하여 원본 TIC(155ms)보다 8.4% 더 빠른 성능을 보였습니다.
• 절제 연구 (Ablation Studies):
  • Token-Choice에서 Expert-Choice 라우팅으로 전환했을 때 BD-Rate가 19.57%에서 15.90%로 개선되었습니다.
  • 공간 TV 정규화를 추가하여 BD-Rate를 15.44%로 더욱 낮추었습니다.
  • GShMLP는 축소된 FFN 확장 비율보다 더 나은 효율성-성능 트레이드오프를 제공했습니다.
  • MoE 레이어를 인코더와 디코더 모두에 배치했을 때 최상의 성능을 보였습니다.

의의 및 주장

본 논문은 MoECodec이 정적인 태스크 특화 커스텀 훈련 패러다임에서 벗어나, 엔드 투 엔드 멀티 태스크 코덱 훈련을 위한 새로운 접근 방식을 제공한다고 주장합니다. 희소하게 활성화되는 동적 계산 할당을 가능하게 함으로써, 이 프레임워크는 단일 모델이 특정 태스크에 중요한 의미적 영역에 계산 자원을 적응적으로 집중할 수 있도록 합니다. 저자들은 이 접근 방식이 인간의 지각(재구성)과 기계 비전 태스크 모두에서 모델 비용(파라미터 및 지연 시간)과 성능 사이의 균형을 성공적으로 맞추었으며, 토큰 수준의 세밀함이 멀티 태스크 이미지 압축을 위한 이질적 계산에 효과적임을 입증했다고 단언합니다.