GoldiMask

Not too much context. Not too little. Just right.

Fine-Tuning Discrete Diffusion LLMs through Submodular Context Selection

GoldiMask chooses which masked tokens to reveal as context by maximizing a non-monotone submodular objective, then weights the remaining targets by how much that context actually helped them. Best average accuracy on every reasoning setting we test, and fewer parallel-decoding iterations at comparable accuracy.

3 backbones3 SFT datasets6 benchmarks2 forward passes, 1 backward
Scroll to open one training step

Stage 0 · one training example

A response the model must learn to reconstruct.

Prompt: Natalia sold clips to 48 of her friends in April, and then half as many in May. How many clips did she sell altogether? The 27 response tokens are the tiles on the right. Hover any tile for its numbers.

Standard SFT hides a random fraction of them and trains on every hidden token equally. GoldiMask makes both choices from measurements instead.

Stage 1 · over-mask

Hide more than the masking rate asks for.

The rate t = 0.45 calls for K = 12 hidden targets. GoldiMask first masks at rate t + ρ, here 0.75, dropping a mask on a larger candidate set. Which candidates stay masked is now a decision, not a coin flip.

K = ⌊tL⌋ = 12  ·  |C| = – candidates  ·  B = |C| − K = – to reveal
visiblemasked candidate

Stage 2 · measure

One forward pass, no gradient.

A measuring light sweeps the grid. Every candidate gets its gold-token probability p, drawn as a pin, and its attention to the other candidates, drawn as threads. Pins glow brightest where the supervision priority peaks.

λ(p) = p(1 − p)³   peaks at p = ¼

A token the model gets right about a quarter of the time is the most valuable target: unsure enough to have room, not so lost that a gradient step cannot move it.

Stage 3 · reveal

Lift the mask where it helps the most.

Revealing a token supports every target that attends to it, but forfeits that token as a target. The objective balances the two. It is submodular but not monotone, so we use maximizers with guarantees; greedy is shown, lifting one mask at a time.

F(R) = Σi ∈ C∖R λ(pi) · φ(Si(R))  ·  φ(s) = s / (s + 0.1)

Greedy reveal order · marginal gain in F

    Stage 4 · weight

    Count each target by how much the context helped.

    A second forward pass on the revealed input gives a new probability for each masked target. Targets whose probability rose, and that still have room, take more of a fixed loss budget. No target may take more than ten times the uniform share.

    ui = [log pi(2) − log pi(1)]+ · λ(pi(2))  ·  wi = min(ui/ν, 10), Σwi = 12

    Targets · p before → after · loss weight

      Stage 5 · after training · inference

      A model that commits more tokens per round.

      At inference a diffusion decoder starts from a fully masked response, predicts every position at once, and commits those above a confidence threshold. GoldiMask trains exactly this prediction problem: a target recovered from partial context. Watch the masks lift in rounds.

      round 0 · committed 0 / 27

      Illustration: confidences and attention are synthetic; the reveal selection and the loss weights are computed live with the paper's rules.

      At a glance

      What one better mask buys.

      4/4
      Reasoning settings where GoldiMask has the best four-benchmark average
      s1K × 3 backbones, LIFT-SFT-12K × LLaDA-8B
      +8.5%
      Largest relative gain in average accuracy over the strongest baseline
      LLaDA-8B on LIFT-SFT-12K, GoldiMask-D vs CART
      16–21%
      Fewer decoding iterations at τ = 0.95 on GSM8K and MATH-500, accuracy within about two points
      all backbones, all three maximizers

      Benchmark accuracy

      Three backbones, three datasets, one recipe.

      All baselines are retrained in the same harness at matched hyperparameters. GoldiMask-D uses deterministic greedy for the reveal set; GoldiMask-R uses randomized greedy, which carries a 1/e guarantee for non-monotone objectives. Values are means over seeds with standard deviations in small type. Bold marks the best entry per column; the line under each GoldiMask row is the relative change against the strongest baseline in that block.

      Parallel decoding · step by step

      Same prompt, same decoder. Watch who finishes first.

      Every board is one fine-tuned model decoding the same response window with a confidence-threshold decoder. Each step, every position above the threshold commits at once and its mask lifts off the text. The clock is shared, so the boards race in real steps. GoldiMask is on top.

      drag to tilt · shift-drag to pan · hover a cell to read its token

      Parallel decoding

      Fewer rounds to the same answer.

      A confidence-threshold decoder commits every unresolved token whose top prediction clears τ, then repeats. Below, the mean number of decoding steps per example for each fine-tuned model, with its accuracy at the finish line. Shorter is faster. Pick any of the twelve settings and any threshold; nothing is hidden.

      0 steps–
      –

      Step counts are iterations of the decoder, each one forward pass over the sequence, so at fixed length they are proportional to compute. They are not wall-clock measurements. On Countdown and Sudoku the picture is mixed; the selector shows those too.