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.
Read original(opens in new tab)