온디바이스 이미지 캡션, 비자기회귀로 12배 빨라졌다!

본문에서는 이미지 캡션 생성 시 자기회귀(Autoregressive) 디코딩의 긴 지연 시간을 해결하기 위해 비자기회귀(Non-Autoregressive) 디코더를 도입했다. 기존 자기회귀 방식은 토큰을 순차적으로 생성하므로, 문장 길이에 비례하여 디코딩 시간이 증가하는 문제가 있었다. 반면, 비자기회귀 방식은 모든 토큰을 병렬적으로 예측하여 시간 복잡도를 O(1)로 줄였다. 특히, 학습 가능한 쿼리 토큰(Learnable Query Tokens)을 활용하여 이미지 인코더의 표현을 쿼리로 사용하고, 병렬적으로 캡션을 생성하는 방식을 채택했다. 이러한 접근 방식은 CVPR 2024의 연구와 ClipCap 논문을 참고하여 구현되었다.

저자는 온디바이스 환경의 제약 속에서 캡션 생성 모델의 성능을 향상시키기 위해 지식 증류(Knowledge Distillation)를 적극 활용했다. 특히, 1단계 지식 증류에서는 자기회귀 모델을 교사 모델로, 비자기회귀 모델을 학생 모델로 설정하여, 학생 모델이 교사 모델의 '눈높이'에 맞춰 학습하도록 유도했다. 2단계에서는 학습된 비자기회귀 모델을 다시 교사 모델로 활용하여, 보다 간결하고 일반적인 스타일의 캡션을 생성하도록 유도했다. 또한, GPT-4o mini를 활용하여 미수락 캡션 데이터를 정제하고, 비문을 수정하는 과정을 통해 최종적으로 수락 비율 89%를 달성했다.

본문에서는 기존 벤치마크 지표(CIDEr, CLIPScore)의 한계를 지적하고, LLM 기반 수락 비율(Accept Ratio)이라는 새로운 평가 지표를 도입했다. 기존 지표들은 모델의 캡션 품질을 정확하게 반영하지 못하는 경우가 많았으며, 특히 반복, 오타, 문법 오류 등을 제대로 감지하지 못했다. 이에 저자는 GPT-4o mini를 활용하여 캡션의 수락/미수락 여부를 판단하고, 이를 기반으로 수락 비율을 측정했다. 이러한 접근 방식은 실제 서비스에서 사용 가능한 캡션의 품질을 객관적으로 평가하고, 모델 개선 방향을 설정하는 데 중요한 역할을 했다.

저자는 온디바이스 환경의 제약 사항을 고려하여 모델 크기를 172MB로 줄이고, 200~400ms의 응답 시간을 달성했다. 이를 위해 이미지 인코더와 텍스트 디코더를 모두 경량화하고, 동적 양자화(Dynamic Quantization)를 적용했다. 특히, 비자기회귀 캡션 생성 방식은 텍스트 디코딩 시간을 30ms로 단축하는 데 기여했다. 또한, 다국어 캡션 생성 시에는 ML Kit을 활용하여 번역 모델을 추가로 다운로드하는 방식을 채택했으며, 모델 다운로드 과정에서의 UX(User Experience)를 고려한 운영 정책을 설계했다.

본 연구에서 제시된 지식 증류 기반 학습 패턴은 다국어 이미지 검색 및 초저지연 이미지 캡션 생성과 같은 서로 다른 문제 해결에 모두 적용되었다. 핵심은 큰 모델의 표현력과 정답 분포를 작은 모델이 학습 가능한 형태로 옮겨오는 것이다. 다국어 이미지 검색에서는 영어 임베딩 공간을 보존하면서 텍스트 인코더를 다국어로 확장했고, 캡션 생성에서는 비자기회귀 방식을 통해 속도를 확보한 후, 캡션 재생성과 다단계 지식 증류를 통해 품질을 향상시켰다. 이러한 접근 방식은 온디바이스 환경에서 모델의 성능과 효율성을 동시에 확보하는 데 효과적인 전략임을 시사한다.