온디바이스 이미지 검색, 지식 증류로 다국어 지원!

by DD
3개월 전
조회수 44

지식 증류(Knowledge Distillation) 기법을 활용하여 모바일 기기 내에서 다국어 이미지 검색 기능을 구현, 5개 언어(영어/일본어/중국어(번체)/태국어/한국어) Recall@5 평균 78% 달성

온디바이스(On-device) 환경의 제약(지연, 프라이버시, 오프라인)을 고려하여 서버가 아닌 모바일 기기 내에서 이미지 이해 기능 구현

PyTorch 모델LiteRT로 변환하여 모델 크기를 70% 감소(205MB → 63MB), 정확도 93.21% 유지, 속도 15.97 it/s 달성

HNSW(Hierarchical Navigable Small World) 기반 벡터 DB를 구축하여 이미지 검색 속도 개선, 10,000장 기준 110ms의 지연 시간 확보

이미지 검색 기능을 비디오 검색으로 확장, CLIP4Clip 모델을 활용하여 파라미터 3배 감소에도 비견할 만한 성능 달성

지식 증류(Knowledge Distillation)를 활용한 다국어 이미지 검색

본문에서는 지식 증류(Knowledge Distillation)를 통해 영어 전용 텍스트 인코더를 다국어로 확장하여 다국어 이미지 검색을 구현했다. 거대 모델(Teacher Model)의 지식을 작은 모델(Student Model)에 전수하여 성능 저하를 최소화하면서 모델 경량화를 달성했다. 특히, 이미지와 텍스트를 같은 임베딩 공간에 매핑하여 언어에 관계없이 의미 기반 검색을 가능하게 했다. 다국어 지원(Multilingual Support)을 위해 5개 언어(영어, 일본어, 중국어(번체), 태국어, 한국어)를 지원하며, Recall@5 평균 78%를 달성했다.

온디바이스(On-device) 환경을 위한 모델 경량화

메신저 환경의 온디바이스(On-device) 이미지 이해 기능 구현을 위해 모델 크기, 속도, 호환성을 모두 고려했다. PyTorch 모델LiteRT로 변환하여 모델 크기를 70% 감소(205MB → 63MB)시키고, 정확도 93.21%를 유지했다. 또한, LiteRT의 연산자 호환성 문제를 해결하기 위해, 에러 함수(erf)를 의사-erf로 대체하고, GroupConvolution을 SeparableConvolution으로 변경하는 등, 다양한 연산 대체 기법을 적용했다. 동적 범위 양자화(Dynamic Range Quantization)를 텍스트 인코더에 적용하여 모델 크기를 추가적으로 줄였다.

LiteRT 변환 과정에서의 기술적 난관

PyTorch 모델을 LiteRT로 변환하는 과정에서 연산자 호환성 문제가 발생했다. LiteRT가 모든 PyTorch 연산을 지원하지 않기 때문이다. 예를 들어, erf(에러 함수)를 의사-erf로 근사 구현하고, GroupConvolution을 SeparableConvolution으로 대체하는 등, 연산 대체(Operator Replacement)를 통해 문제를 해결했다. 또한, LiteRT 버전 정합성 문제를 고려하여, 모바일 앱에 탑재된 LiteRT 런타임 버전과 호환되는 모델을 생성하기 위해, 모델 변환 단계에서부터 모바일에서 실제로 가능한 연산만 사용하도록 제약을 걸었다. ONNX2TF를 사용하여 모델 변환을 수행했다.

HNSW 기반 벡터 DB 구축 및 튜닝

온디바이스 이미지 검색을 위해 HNSW(Hierarchical Navigable Small World) 기반의 벡터 DB를 구축했다. HNSW는 검색 속도와 재현율(Recall)의 균형을 잘 잡을 수 있는 ANN(Approximate Nearest Neighbor) 인덱스 구조이다. HNSW 하이퍼파라미터 탐색을 통해, 데이터 크기(3K/10K/30K)와 파라미터(neighbors_per_node, indexing_search_count)를 그리드로 탐색하여, 검색 지연과 재현율 간의 파레토 최적점을 찾았다. 30K 이미지 기준 M=128, ef_construction=100으로 설정하여, 10,000장 기준 8~50ms(평균 10ms)의 검색 지연 시간을 확보했다. 벡터 적재 배치(Vector Loading Batch)는 1000장 단위로 DB 삽입했다.

전처리/후처리 및 토큰화기 구현

모바일 환경에서 모델의 성능을 최적화하기 위해, 전처리/후처리 과정에 많은 노력을 기울였다. 특히, 학습 환경과 동일한 전처리 파이프라인을 구현하는 것이 중요했다. 이미지 전처리 시, RGB 순서를 확인하고, 크기 조정 방식을 명시적으로 지정하여 임베딩 결과의 일관성을 확보했다. 또한, 코사인 유사도 계산을 위한 L2 정규화(L2 Normalization)를 수행했다. 텍스트 입력을 위한 토큰화기(Tokenizer)를 직접 구현하여, 다국어 문자 처리, 대소문자 처리, 특수 문자 처리 등, 다국어 환경에 맞는 토큰화를 수행했다. ByteBuffer 할당 및 복사 최소화를 통해 추론 성능을 개선했다.

메신저용 온디바이스 이미지 모델 학습기 1편: 지식 증류로 확장한 다국어 이미지 검색