Day 2 • 10:50 R0
Learn Japanese with Manga! Building a Multimodal Search Engine with Gemini Embedding 2
Machine LearningChinese (English slides)
@x0mg
Let's Build a Nutrition Agent with JAX: RL Fine-Tuning a Small Model
John Lu
Day 2 • Sun, Oct 18
14:00 - 14:45
Location
R1
Language
English
Category • Level
Machine Learning • Expert
Snap a photo of your lunch and have an app log it ("Lentil Curry, ~320 kcal, 34g protein"). The obvious build throws a giant multimodal model at it, but that's slow and costly to run all day.
We build this coach as a multimodal agent: a vision-language model identifies food, calls a nutrition database, and returns a structured log. Then comes the real problem—making it cheap. We reinforcement-learn a small model to tool-call as reliably as one ten times its size, using JAX-based tooling (Tunix for GRPO) on a free Colab TPU.
You'll meet the three JAX concepts behind the stack (grad, jit, vmap), understand why its hardware-accelerated functional design is ideal for RL, and see how reward functions are designed for tool-calling. Watch a 4B-class model (e.g., PaliGemma) close the gap on a 31B baseline (e.g., LLaVA) across strict format metrics. You'll leave knowing when to RL-fine-tune a small model and when it breaks. All code is open-source.
Description
We build up with runnable code and a public repo attendees can use afterward.
1. The agent system architecture. We frame the product—photo in, accurate nutrition log out. We map out the broader workflow: a user uploads an image → the Vision-Language Model (VLM) analyzes it and generates a JSON-formatted tool call → the Python backend executes the nutrition API lookup → the VLM synthesizes the final structured log. This grounds the audience in how the components integrate before we optimize them.
2. Why the obvious build is expensive. Running a large multimodal model (like a 31B parameter LLaVA/Idefics) for every meal is slow and costly. We make the cost concrete and set up the central question: can a much smaller 4B-class model (like PaliGemma or Qwen-VL) do the tool-calling job just as well?
3. The JAX you actually need (and why JAX). A short, hands-on tour of grad, jit, and vmap. Crucially, we explain the architectural rationale for JAX: why its functional purity and XLA compiler are perfectly suited for RL fine-tuning on TPUs. We provide an honest map of where JAX sits beneath the model tooling (Flax for structure, Optax for optimization, Tunix for orchestration).
4. Shrinking the model with RL. The heart of the talk: GRPO-style reinforcement fine-tuning. We explain the actor/reference setup, then design a reward function that scores the things that matter. Instead of static labels, we score based on Format Adherence (did it output valid JSON?) and Call Success (did the API return a valid hit?). We orchestrate the run with Tunix on JAX, stream data with Grain (Food101 images), and show the reward code live.
5. Evaluation Setup & Results. We define our strict evaluation metrics: Tool Call Success Rate, Format Error Rate, and Caloric Deviation. Using a held-out evaluation set of complex meals, we compare the small fine-tuned 4B model against the 31B baseline. We demo the model running end-to-end on a single accelerator (free Colab/Kaggle TPU).
6. Takeaways and Failure Cases. When does this approach win? We summarize the transferable JAX/Tunix patterns. Importantly, we highlight key failure cases (e.g., when the small model hallucinates obscure ingredients, or fails on highly obscured/blended foods like smoothies) so the audience can honestly evaluate the limits of small-model RL.
John Lu
John is a Senior Machine Learning Engineer from a Top-tier company, with hands-on, production experience building Large Language Models, Vision Language Models, and RAG and tool-calling systems. He likes pulling apart the "magic" of modern AI and rebuilding it in clear, runnable code, and applying it to problems he lives with: optimizing his own diet and tracking nutrient-dense foods is a long-standing habit, which makes this project something he actually uses. He has spoken at PyCon APAC and MOPCON on applied ML and on-device AI, including TasteGram (image assessment with TensorFlow & Python) and Kimoji (an On-Device LLM mood-tracker system).
