Density-Aware Reward Scaling

In one sentence: when outputs are images or free text, no two are ever exactly the same, so “how often did this come up?” becomes “how many near-copies are there?”. DARS measures that with a simple similarity count and divides the reward by it.

Background

RL fine-tuning tends to make generative models repetitive: the same style of image, the same kind of reply. Our IPS work traced this to the training objective. It rewards outcomes in proportion to how often they already occur, so popular outcomes snowball. The fix was to divide each reward by the outcome’s frequency.

That requires spotting when two outputs are the same. For images and paragraphs, two outputs are almost never exactly identical, even when they are near-copies: two moody night-time snowmen with slightly different pixels, or two replies that differ by one word. Exact counting sees every output as unique and does nothing.

The idea: count near-copies

Density-Aware Reward Scaling (DARS) replaces exact counting with a soft count:

Map each output in the group to an embedding, a vector in which similar things sit close together (we use DINOv2 for images and a sentence-embedding model for text).
For each output, add up how close every other output in the same group is to it, using a Gaussian kernel. Exact duplicates count as 1, near-copies as almost 1, and very different outputs as roughly 0. This sum is the output’s local density.
Divide the reward by the density, then run the normal GRPO update.

Figure 1 · Counting near-copies instead of exact copies

Each generated output is mapped to an embedding. Its density is the sum of a Gaussian kernel over the other outputs in the same group, so near-copies count as partial repeats. Dividing the reward by this density takes most of the extra weight away from the crowd. The numbers here come from the kernel applied to the points shown.

That is the whole method. There is no learned density model, critic, extra loss or architecture change. It is a few lines of reward processing before the usual advantage calculation.

A useful side effect is that the reward controls what kind of diversity you get. The ideal end point is a model that samples each output in proportion to its reward. You don’t have to define “novelty” separately. You describe what is good, and DARS spreads the model’s attention across everything that is good.

Does it really sample in proportion to reward?

To check, we built a test with a known right answer. An agent moves across a square in small continuous steps and decides where to stop. The reward has four equally tall peaks, so a perfect sampler stops near each peak 25% of the time.

Figure 2 · Continuous navigation with four equal peaks

An agent moves across a square and stops somewhere. The reward has four equally tall peaks, so the ideal policy stops near each one a quarter of the time. GRPO always goes to the same peak. DARS-GRPO splits its samples almost exactly evenly, and on this measure it is even closer to the target than a GFlowNet (Table 4 of the paper, averaged over 5 seeds).

GRPO puts every sample on one peak. DARS-GRPO, the same algorithm with the same network plus density scaling, splits its samples almost evenly. On measures of the overall shape of the distribution, it is even closer to the target than a GFlowNet, which is designed for this kind of sampling.

Diverse images from prompt rewriting

A practical use: a small language model learns to rewrite a user’s prompt so that a frozen Stable Diffusion model produces images that are faithful to the request and look good. With plain GRPO, the rewriter finds one trick and repeats it. For “A snowman standing in the yard”, every rewrite became a variation of “…moody evening”, and every picture was a similar night scene. With DARS the rewrites spread out: “…details of Paris”, “…cherry blossoms”, “…evening light”.

Figure 3 · DARS for text-to-image prompt rewriting

Only the small language model is trained. The image model and both scoring models stay frozen. Density is measured on the images, not on the text, because the goal is visually different pictures and not just reworded prompts.

Across two datasets (MS-COCO and Lexica), two rewriter models (Qwen3-0.6B and GPT-2) and two reward definitions, DARS-GRPO gave the highest image diversity of all the prompt-rewriting methods we tested. Image quality and prompt faithfulness stayed close to the baselines.

Image diversity (Vendi score, MS-COCO, alignment reward)	GRPO	DARS-GRPO
Qwen3-0.6B rewriter	9.5	20.1
GPT-2 rewriter	8.8	22.9

The one setting that needs care is the kernel width, which decides what counts as “nearby”:

Figure 4 · Choosing the neighbourhood size

The kernel width decides what counts as "nearby". If it is too small, nothing has neighbours and DARS changes nothing. If it is too large, everything is everyone's neighbour and the scaling cancels out. Diversity peaks in between (Table 3 of the paper, 5 seeds).

Open-ended text

We also fine-tuned a 2-billion-parameter language model on real user prompts and tested it on prompts that have many good answers. DARS roughly doubled the number of functionally different answers among 16 samples (12.0 versus 5.9 for GRPO), and also doubled semantic diversity. This came with a trade-off. The average reward-model score was lower, because the model spreads its probability over more good answers instead of polishing one. The usefulness of answers that are both good and new still went up.

Limits

DARS is only as good as its embedding. If distances in that space don’t reflect what people consider “different”, the scaling won’t either. The pairwise comparison grows with the square of the group size. And our evidence for reward-proportional sampling is empirical, not a general guarantee.