BiGain: Unified Token Compression for Joint Generation and Classification

이 논문은 생성 품질을 유지하면서 분류 성능까지 향상시키는 훈련 불필요의 플러그인 프레임워크인 'BiGain'을 제안하며, 주파수 분리 원리를 기반으로 한 토큰 병합 및 KV 다운샘플링 연산자를 통해 확산 모델의 가속화 시 생성과 분류의 균형을 최적화합니다.

핵심

🎨 비유: "화실의 붓질 정리하기"

생각해 보세요. 한 화가가 거대한 캔버스에 아주 정교한 그림을 그리고 있습니다.

• 기존 방식 (기존 가속화 기술): 화가가 그림을 더 빨리 그리기 위해 캔버스 전체를 무작위로 잘라내거나, 비슷한 색을 가진 부분끼리 뭉개버립니다.

  • 결과: 그림은 여전히 예쁘게 보일 수 있지만, 고양이의 귀 끝이나 나뭇잎의 섬세한 줄무늬 같은 중요한 디테일이 사라져 버립니다. 그래서 "이게 고양이야, 강아지야?"를 구분하는 AI 는 엉뚱한 답을 내놓게 됩니다. (그림은 예쁘지만, 분류는 망함)

• BiGain 의 방식 (이 논문의 제안): 화가가 그림을 그릴 때, **"어디는 간결하게 줄이고, 어디는 디테일을 꼭 지켜야 할지"**를 아주 똑똑하게 판단합니다.

  • 핵심 아이디어: 그림을 '주요한 내용 (배경, 형태)'과 '세부적인 디테일 (가장자리, 질감)'로 나누어 생각합니다.

🔍 BiGain 의 두 가지 마법 도구

BiGain 은 그림을 빠르게 만들면서 중요한 정보를 잃지 않게 해주는 두 가지 도구를 사용합니다.

1. 라플라시안 게이트 (Laplacian-gated Token Merging)

비유: "잔잔한 호수와 거친 파도"

• 원리: 그림의 각 부분을 살펴봅니다.
  • 잔잔한 호수 (매끄러운 부분): 하늘이나 벽처럼 색이 일정하고 복잡한 게 없는 곳은 서로 합쳐서 (Merge) 공간을 줄입니다.
  • 거친 파도 (세부적인 부분): 고양이의 눈, 나뭇가지, 옷의 주름처럼 가장자리와 질감이 뚜렷한 곳은 절대 합치지 않고 그대로 둡니다.
• 효과: 불필요한 공간은 줄여서 속도를 높이지만, 사물을 구분하는 데 필수적인 '가장자리'는 그대로 보존합니다. 그래서 그림이 예쁜 것은 물론, AI 가 "이건 고양이구나!"라고 정확히 맞힐 수 있습니다.

2. KV 다운샘플링 (Interpolate-Extrapolate KV Downsampling)

비유: "지도 읽는 방법"

• 원리: AI 가 그림을 그릴 때 '무엇을 그릴지 (질문, Query)'와 '어디에 그릴지 (키/값, Key/Value)'를 결정합니다.
  • 질문 (Query): "이곳에 고양이를 그려야 해?"라고 묻는 역할입니다. 이 부분은 원본 그대로 남겨둡니다. (정확한 위치 파악을 위해)
  • 정보 (Key/Value): "고양이는 이런 모양이야"라는 정보입니다. 이 정보는 적당히 줄여서 효율을 높입니다.
• 효과: AI 가 "어디에 집중해야 할지"는 정확히 알면서, 기억해야 할 정보량은 줄여서 속도를 냅니다. 마치 지도를 볼 때 "내가 어디에 있는지"는 정확히 알면서, "주변의 모든 나무 이름"까지 다 외울 필요는 없게 만드는 것과 같습니다.

📊 왜 이것이 중요한가요?

기존 기술들은 "그림을 빨리 그리는 것"에만 집중해서, 그림을 그리는 AI 가 동시에 '감별사' 역할도 할 수 있다는 점을 무시했습니다.

• 기존: 그림은 빨리 그려지지만, 분류 정확도가 뚝 떨어집니다. (예: 70% 속도 향상 시, 분류 정확도 20% 하락)
• BiGain: 그림도 예쁘게, 분류도 정확하게!
  • 실험 결과, 그림을 그리는 속도 (계산량) 를 70% 줄였을 때, 분류 정확도는 오히려 7% 이상 향상되었고, 그림의 질도 떨어지지 않았습니다.

💡 결론: "한 마리 토끼 두 마리 잡기"

BiGain 은 **"그림을 잘 그리는 능력 (생성)"**과 **"그림을 잘 구별하는 능력 (분류)"**이라는 두 마리 토끼를 동시에 잡는 기술입니다.

이 기술은 AI 모델을 재학습 (Training) 시킬 필요 없이, 기존 모델에 바로 끼워 넣을 수 있어 (Plug-and-play) 비용 절감과 효율성 측면에서 매우 유용합니다. 앞으로 의료 영상 진단, 산업용 불량 검사, 보안 감시 등 **"그림을 보고 판단해야 하는 모든 분야"**에서 AI 를 더 빠르고 정확하게 사용할 수 있는 길을 열어줍니다.

한 줄 요약:

"BiGain 은 AI 가 그림을 그릴 때, 중요한 디테일은 살리고 불필요한 부분은 잘라내어, 그림도 더 빨리 그리고 사물도 더 잘 구분하게 해주는 똑똑한 '붓질 정리' 기술입니다."

기술 요약

1. 문제 정의 (Problem)

확산 모델 (Diffusion Models) 은 현대 생성 시스템의 핵심이지만, 샘플링 과정에서의 높은 계산 비용으로 인해 토큰 병합 (Token Merging) 이나 다운샘플링 (Downsampling) 과 같은 가속화 기법들이 활발히 연구되고 있습니다. 그러나 기존 방법들은 주로 생성 품질 (Synthesis Quality, 예: FID) 을 최적화하는 데 초점을 맞추고 있어, 분별 능력 (Discriminative Capacity, 예: 분류 정확도) 을 무시하는 경향이 있습니다.

실제로 확산 모델은 생성뿐만 아니라 선형 프로브 (Linear Probe), 특징 증류, 확산 기반 분류 프로토콜 등을 통해 이미지 분류 작업에도 널리 활용됩니다. 저자들은 기존 가속화 기법들이 생성에는 큰 손상을 주지 않더라도, 분류 성능을 급격히 저하시킨다는 사실을 발견했습니다. 이는 에지 (Edge), 질감 (Texture), 고대비 경계와 같은 고주파수 (High-frequency) 세부 정보가 분류에 결정적인 역할을 하는데, 단순한 토큰 압축이 이러한 정보를 과도하게 제거하기 때문입니다. 즉, "잘 보이는 (생성)" 것과 "잘 분류되는 (분별)" 것 사이의 간극을 해소할 수 있는 새로운 압축 원칙이 필요했습니다.

2. 방법론 (Methodology)

저자들은 토큰 압축을 생성 충실도 (Generative Fidelity) 와 분별 유용성 (Discriminative Utility) 을 동시에 보존하는 이중 목적 (Bi-objective) 최적화 문제로 재정의했습니다. 핵심 아이디어는 주파수 분리 (Frequency Separation) 입니다.

• 핵심 원리: 중간 특징 (Intermediate features) 의 신호를 주파수 인식 표현으로 매핑하여, 에지와 미세 질감 같은 고주파수 세부 정보와 전체적인 의미 (Global Semantics) 를 담은 저/중주파수 콘텐츠를 분리합니다.
• 설계 원칙: 분류에 필수적인 고주파수 성분을 보존하면서, 생성의 일관성을 유지하는 저/중주파수 대역도 유지하는 균형 잡힌 스펙트럼 보존 (Balanced Spectral Retention) 을 목표로 합니다.

이를 구현하기 위해 BiGain이라는 훈련 불필요 (Training-free), 플러그 앤 플레이 (Plug-and-play) 프레임워크를 제안하며, 두 가지 주파수 인식 연산자로 구성됩니다.

A. 라플라시안 게이트드 토큰 병합 (Laplacian-Gated Token Merging, L-GTM)

• 목적: 스펙트럼적으로 매끄러운 (Low-frequency) 토큰끼리는 병합하되, 고대비 (High-contrast) 세부 정보를 가진 토큰의 병합은 억제합니다.
• 작동 방식:
  1. 숨겨진 상태 (Hidden state) 토큰에 라플라시안 필터 (Laplacian Filter) 를 적용하여 국소 주파수 점수를 계산합니다.
  2. 각 그리드에서 점수가 가장 낮은 (매끄러운) 토큰을 '목적지 (Destination)'로 선택하고, 나머지를 '출처 (Source)'로 설정합니다.
  3. 출처와 목적지 토큰 간의 유사도 (Cosine Similarity) 를 기반으로 최상위 쌍을 선택하여 평균화 (Merge) 합니다.
• 효과: 에지와 질감 같은 미세 구조를 보존하면서 중복된 평탄한 영역만 압축하여 계산 비용을 절감합니다.

B. 인터폴레이트 - 엑스트라폴레이트 KV 다운샘플링 (Interpolate-Extrapolate KV-Downsampling, IE-KVD)

• 목적: 어텐션 (Attention) 계산을 줄이기 위해 키 (Key) 와 값 (Value) 을 다운샘플링하되, 쿼리 (Query) 는 원본 해상도로 유지합니다.
• 작동 방식:
  • 쿼리 (Q) 는 그대로 유지하여 모델의 정밀한 위치 인식 및 어텐션 능력을 보존합니다.
  • 키 (K) 와 값 (V) 은 최근접 (Nearest) 풀링과 평균 (Average) 풀링 사이의 가변적인 보간/외삽 (Interpolate-Extrapolate) 을 통해 다운샘플링합니다.
  • 제어 가능한 매개변수 ($\alpha$) 를 통해 주파수 보존과 계산 효율 사이의 균형을 조절합니다.
• 효과: 메모리와 FLOPs 를 줄이면서도 분류에 중요한 토큰별 어텐션 정밀도를 유지합니다.

3. 주요 기여 (Key Contributions)

1. 이중 목적 토큰 압축 프레임워크: 확산 모델의 토큰 압축을 생성과 분류를 동시에 고려하는 문제로 재정의하고, 훈련 없이 적용 가능한 실용적인 솔루션 (BiGain) 을 제시했습니다.
2. 최초의 통합 연구: 가속화된 확산 모델 하에서 생성과 분류를 동시에 연구하고 개선한 최초의 프레임워크입니다.
3. 주파수 인식 설계 가이드: "신호가 매끄러운 곳에서는 병합하고, K/V 는 주파수 균형을 고려하여 다운샘플링하되 Q 는 보존한다"는 실용적인 설계 원칙을 제시하여, 향후 배포 가능한 이중 목적 생성 모델 개발에 기여합니다.

4. 실험 결과 (Results)

저자들은 DiT 및 U-Net 기반 백본과 ImageNet-1K, ImageNet-100, Oxford-IIIT Pets, COCO-2017 등 다양한 데이터셋에서 실험을 수행했습니다.

• 분류 성능 향상: 기존 가속화 방법 (ToMe, ToDo 등) 은 분류 정확도가 급격히 떨어지는 반면, BiGain 은 가속화 비율이 높아질수록 기존 방법 대비 분류 정확도를 크게 개선했습니다.
  • 예시: Stable Diffusion 2.0 에서 70% 토큰 병합 시, ImageNet-1K 분류 정확도가 7.15% 향상되었습니다.
• 생성 품질 유지 및 개선: 분류 성능을 개선하면서도 생성 품질 (FID) 을 유지하거나 오히려 개선했습니다.
  • 예시: 위 조건에서 FID 는 0.34(1.85%) 개선되었습니다.
• 균형 잡힌 스펙트럼의 중요성: 고주파수 세부 정보와 저/중주파수 의미 내용을 모두 보존하는 것이 확산 모델의 토큰 압축에 대한 신뢰할 수 있는 설계 규칙임을 입증했습니다.
• Ablation Study: 라플라시안 게이팅을 제거하거나 주파수 균형을 무시한 다운샘플링을 적용할 경우 분류 성능이 크게 저하됨을 확인하여, 제안된 주파수 인식 메커니즘의 필수성을 입증했습니다.

5. 의의 및 결론 (Significance)

BiGain 은 단일 확산 백본이 생성과 분류라는 두 가지 상반된 작업을 동시에 수행해야 하는 현실적인 요구 (의료 영상, 안전 필수 감지, 산업 검사 등) 에 부응합니다. 기존 가속화 기법들이 "생성만 잘되면 된다"는 편향에서 벗어나, 생성과 분별 능력을 모두 보존하는 균형 잡힌 토큰 압축의 중요성을 강조했습니다.

이 연구는 추가적인 훈련 없이 기존 확산 모델에 바로 적용 가능하여, 저비용으로 배포 가능한 이중 목적 (Dual-purpose) 생성 시스템의 실현을 가능하게 하며, 확산 모델의 효율성과 활용 범위를 확장하는 중요한 이정표가 됩니다.