On-Device Image Model Training for Mess

This post describes an on-device image captioning system for mobile messenger apps. Because autoregressive vision-language models took more than five seconds to generate captions, the team replaced them with a non-autoregressive decoder, reducing latency to roughly 200–400 ms. They then used LLM-based acceptance evaluation, caption re-generation, and multi-stage knowledge distillation to improve quality while keeping the model at 172 MB.

Why Conventional Captioning Was Unsuitable

• Models such as BLIP-2, MobileVLM, PaliGemma, and MiniCPM were too large or slow for mobile deployment.
• BLIP-1 was selected as a practical baseline because of its smaller size and clear licensing, but still required more than five seconds after quantization.
• Autoregressive decoding generates tokens sequentially, requiring one decoder pass per token.
• On a Samsung Galaxy Fold 4, the initial model required about 142 ms per token, or approximately 2.8 seconds for 20 tokens.
• Mobile UX required stable latency in the hundreds of milliseconds, including cold-start and variable-device conditions, so simple model compression was insufficient.

Non-Autoregressive Caption Generation

• The system predicts all caption tokens in parallel using a fixed set of learnable query tokens.
• This changes the decoding cost from roughly O(T) for autoregressive generation to near O(1) through parallel processing.
• The architecture consists of:
  • An image encoder reused from the previous system
  • Image embeddings injected as a prefix, following the ClipCap approach
  • A 66.4-million-parameter Transformer-based text decoder
  • Twenty learnable query tokens for short captions
• Query-CTC loss addresses the alignment problem caused by predicting tokens simultaneously.
• The resulting model generated captions in about 200 ms, achieving the required speed improvement.

Speed Improved, but Caption Quality Declined

• Standard metrics such as CIDEr and CLIPScore appeared acceptable.
• Manual inspection revealed frequent:
  • Repeated words, such as “a desk with a computer on a desk”
  • Spelling errors, such as “a people ons”
  • Grammatical problems
  • Incomplete captions, such as “a”
• These defects made the model unsuitable for direct use in a messaging product.

LLM-Based Acceptance Evaluation

• The team introduced an “accept ratio” based on GPT-4o mini judgments.
• Captions were classified as either accept or non-accept.
• The evaluation checked for:
  • Duplicate content
  • Errors
  • Clarity and grammatical correctness
• This better reflected production usability than conventional image-captioning benchmarks.
• The low acceptance rate confirmed that CIDEr and CLIPScore alone could not measure whether captions were appropriate for users.

Data Quality and Knowledge Distillation

• Analysis showed that the training data contained inconsistent and noisy captions:
  • A mixture of very short and overly long descriptions
  • Unnecessary OCR-like attempts to describe text in images
  • Uneven language quality
• The smaller 66.4-million-parameter model also had less representational capacity than BLIP-1’s 110 million parameters.
• Generating an entire sentence in one pass was especially difficult for the compact non-autoregressive model.
• The training pipeline was redesigned as an iterative quality-improvement loop:
  • Train a baseline using the original data
  • Identify failures with LLM-based acceptance evaluation
  • Re-caption poor-quality training examples
  • Distill knowledge from a larger teacher model into the student model
  • Replace or refine rejected samples and repeat
• Architectural scaling and metric optimization did not consistently improve acceptance rates, while re-captioning and knowledge distillation produced more meaningful gains.

The practical recommendation is to design on-device captioning around the actual product experience: prioritize parallel decoding for latency, measure quality with production-oriented acceptance criteria, and use carefully curated data plus knowledge distillation to make compact models reliable.