On-Device Image Model (opens in new tab)
The post describes building an on-device image understanding system for messaging apps, with semantic image search as the first focus. Its central strategy was knowledge distillation: a multilingual student text encoder learned to reproduce the embedding space of a strong, English-only teacher model. This preserved most English retrieval quality while enabling Japanese, Traditional Chinese, Thai, and Korean search, achieving an average Recall@5 above 78%. ## Why Messaging Apps Need On-Device Image Understanding - Images are often treated simply as “a photo,” unlike text messages, which can support search, summaries, and notification previews. - Image understanding could improve: - Notifications: “Sent one photo” → “Sent a photo of a dog” - Search: queries such as “dog,” “puppy,” or “a cat inside a box” - Recommendations: automatic image classification and organization - Shared image-text embeddings allow semantically equivalent phrases—such as “dog,” “puppy,” and “개”—to retrieve the same images. ## Why the Model Had to Run On-Device - **Latency:** Network round trips make notifications and search less predictable and responsive. - **Privacy:** Sending photos, captions, or embeddings to a server increases privacy risks. - **Offline support:** The feature should work in subways, airplanes, roaming environments, and unstable networks. - **Mobile constraints:** The model needed to run on both Android and iOS with limited memory and compute resources. - The project targeted a model under **200 MB**, response times within a few hundred milliseconds including cold start, and LiteRT compatibility. ## Project Goals and Evaluation - The image search system needed to: - Retrieve images by semantic meaning rather than keyword matching. - Support English, Japanese, Traditional Chinese, Thai, and Korean. - A separate captioning system was designed to generate short, natural descriptions of roughly eight words or fewer. - Search quality was measured using: - Image-to-Text Recall@5 - Text-to-Image Recall@5 - Caption quality was evaluated with CIDEr, CLIPScore, and an LLM-based acceptance ratio designed to detect repetition, typos, and grammatical problems. ## Why Translation Was Not Enough The initial approach translated each query into English before using an English-only image-text model: ```text Query → Language detection → Translation → English text encoder → Embedding → Search ``` This approach introduced several problems: - **Quality loss:** Informal terms or short queries could be mistranslated. For example, “멍멍이” might be interpreted as “barking” instead of “dog.” - **Additional latency:** Translation adds a fixed cost before text encoding. - **Inconsistent results:** Translation quality varies by language pair and wording. - **Operational complexity:** Each additional language requires more models, updates, and failure handling. Training a multilingual image-text model from scratch would require substantial data and compute. Instead, the project retained the proven English image embedding space and expanded only the text encoder. ## Knowledge Distillation for Multilingual Search - The original English text encoder served as the frozen **teacher**. - A copied text encoder served as the trainable **student**. - English text was passed to the teacher, while corresponding multilingual text was passed to the student. - The student was trained to match the teacher’s embeddings using mean squared error (MSE). ```text teacher_embedding = teacher(English text) student_embedding = student(Multilingual text) loss = MSE(teacher_embedding, student_embedding) ``` The image encoder remained frozen so that the established image-text embedding space would not be disrupted. Important implementation considerations included: - Ensuring the tokenizer handled multilingual characters correctly. - Defining consistent case-insensitivity rules. - Balancing training samples across languages. - Matching training-time preprocessing and tokenization with mobile inference behavior. ## Retrieval Results - English performance declined slightly: - Image-to-Text Recall@5: **79.58% → 76.56%** - Text-to-Image Recall@5: **75.89% → 74.47%** - Multilingual performance improved from below **10% average Recall@5** to above **78%**, roughly a sevenfold improvement. - Japanese achieved **81.94%**, exceeding the original English model in the reported evaluation. - Traditional Chinese, Thai, and Korean also reached practically usable retrieval quality. The trade-off—slightly lower English performance in exchange for four additional languages—provided substantially greater overall product value. ## Converting the Model to LiteRT - LiteRT was selected because it officially supports both Android and iOS and provides mobile-oriented operators, quantization, and optimization tools. - Core ML was rejected because it is iOS-specific and introduced conversion and long-term cross-platform maintenance concerns. - Conversion required addressing unsupported PyTorch operators. - For example, LiteRT did not support `erf`, so the model’s implementation had to replace it with a compatible pseudo-`erf` operation. The resulting approach demonstrates that knowledge distillation can efficiently extend an existing English image-text model to multiple languages while preserving its on-device deployment advantages.