techlist.io
KREN
line

On-Device Image Model (opens in new tab)

pytorchknowledge-distillationlitertembedding-modelcosine-similarityimage-searchmultilingual-supportond-device

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:

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).
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.