Tuning computer vision models with task rewards

TL;DR

The paper tackles misalignment between computer vision models and their intended usage by introducing a simple two-stage approach: pretrain with maximum likelihood and then fine-tune with task-reward optimization using Reinforce. This direct reward tuning is demonstrated across panoptic segmentation, object detection, colorization, and image captioning, improving task-aligned metrics beyond standard MLE performance. Key contributions include task-specific reward designs (e.g., PQ, recall, mAP, colorfulness, CIDEr) and empirical evidence that reward-driven tuning yields substantial gains while maintaining alignment with the intended usage. The results suggest reward optimization as a practical, general mechanism to control nontrivial task risks in vision systems, with potential extensions to human feedback and more complex rewards.

Abstract

Misalignment between model predictions and intended usage can be detrimental for the deployment of computer vision models. The issue is exacerbated when the task involves complex structured outputs, as it becomes harder to design procedures which address this misalignment. In natural language processing, this is often addressed using reinforcement learning techniques that align models with a task reward. We adopt this approach and show its surprising effectiveness across multiple computer vision tasks, such as object detection, panoptic segmentation, colorization and image captioning. We believe this approach has the potential to be widely useful for better aligning models with a diverse range of computer vision tasks.

Figures (6)

• Figure 1: By tuning a strong, pretrained model with a reward that relates to the task, we can significantly improve the model's alignment with the intended usage.
• Figure 2: Step 1: Maximum-likelihood training. In a first step, the model is trained to maximize the likelihood of the ground-truth annotations. This is the most common way to train a model and corresponds to learning to imitate the collected data.
• Figure 3: Step 2: Reward tuning. In a second step, the model is further trained to maximize a reward function. This is done using Reinforce by adjusting the likelihood of model outputs according to their reward.
• Figure 4: Random examples demonstrating how UViM colorization model kolesnikov2022uvim predictions change after tuning with a "colorfulness" reward. See text in \ref{['sec:color']} for the details on the reward function.
• Figure 5: Analysis of reward distribution before and after tuning the model for the image captioning task. Measured as mean of 1024 validation examples. In quantile plot (a) the AUC shows the difference between the methods while highlighting that after reward optimization the chance of sampling low-performing samples is greatly reduced. Plot (b) the max reward out of N samples shows that the MLE model includes high-quality outputs in a large enough pool. Plot (c) it is not possible to identify the best outputs using likelihood as observed by the low reward when using the most likely out of N samples even when using $10000$ samples.