Joint Supervision and Balanced Batching for Spatial Decisions

Why People Wanted It

A spatial answer may not be one decision. A numerical estimate can be refined through successive choices, and a point prediction can make decisions around image crops. Training each stage with its own backbone pass repeats work; combining examples of very different lengths can waste computation in padding or leave some GPUs waiting. The real question is how to teach the full decision path without paying for a separate pass at every step—and how to keep the resulting workload balanced.

The Core Move

Put the known refinement path into one causal sequence, score every supervised decision in that sequence, and divide each question’s total loss across its stages:

$L_q = \frac{1}{S_q}\sum_{s=1}^{S_q} \ell_{q,s}$

Here, $S_q$ is the number of supervised decisions for question $q$, and $\ell_{q,s}$ is its categorical cross-entropy at stage $s$.

Quiz

A question has three supervised refinement stages. Why divide its summed stage losses by three before averaging across questions?

To keep a multi-stage question from receiving more total weight solely because it has more stages.

To let earlier stages attend to later target choices.

To make inference evaluate all dependent stages in parallel.

To remove the need for categorical cross-entropy at each decision.

$ check

Why Teacher Forcing Can Supervise the Whole Path

Now the version for readers who want more of the machinery: the key constraint is causal order, not that each stage gets a separate forward pass. The annotated answer determines the complete target path and any crops before training. Serialize each decision’s target choice—and any subsequent crop—after that decision position. A causal model can then compute every stage loss in one backbone pass: earlier positions cannot attend to future target choices or crops, while teacher forcing supplies the known preceding choices.

This makes training different from inference. At inference, the next candidates and any crop depend on the model’s current selection, so stages must be evaluated sequentially. Joint supervision reuses a training pass; it does not make a dependent inference path parallel.

Why Stage Averaging Preserves Question Weight

Here is the part that gets technical: summing stage losses would give questions with longer refinement paths more total influence. Instead, average the categorical cross-entropies within each question, then average over the real questions in the optimizer batch. A one-stage choice and a three-stage point example therefore each contribute one question’s total weight, rather than the point example counting three times as much.

The batch denominator is the number of actual questions. Any tail placeholders used to make batch shapes line up receive zero weight; they do not change the normalization. This defines the intended objective even when the distributed batch contains padding placeholders.

Why Similar-Length Work Still Needs Global Scheduling

Now the deeper systems detail: examples differ not only in token length, but also in visual-token work. Point examples add two image crops, and choice and numerical examples can differ in candidate length and stage count. Sort expanded sequences by length within each global optimizer batch to form similar-length microbatches, then assign the most expensive remaining microbatch to the GPU with the least assigned work. The cost accounts for padded sequence length and visual-token work.

The scheduling changes where work runs, not which questions belong to the optimizer batch or their loss weights. Keep each QA as an independent causal row: packing unrelated questions together would require isolating the model’s Gated DeltaNet states as well. The open question is how much throughput the schedule recovers on a particular hardware and data mix.

Minimal Causal Code

This toy loop shows the per-question stage average and zero-weight placeholders. It omits the model-specific serialization, crop handling, and distributed scheduler.

import torch
import torch.nn.functional as F

# Each row is one question; valid stages are already in causal order.
# logits: [questions, stages, classes]
# targets: [questions, stages]
# valid: [questions, stages], False for absent stages/placeholders
logits = model(input_ids)
per_stage = F.cross_entropy(
    logits.transpose(1, 2), targets, reduction="none"
)

stage_counts = valid.sum(dim=1).clamp_min(1)
per_question = (per_stage * valid).sum(dim=1) / stage_counts
real_question = stage_counts.gt(0)  # production code tracks tail placeholders explicitly
loss = per_question[real_question].mean()
loss.backward()

What You Gain, and What You Lose

Joint supervision avoids a separate language-backbone training pass for every refinement stage, and stage averaging keeps a longer path from automatically receiving more total question weight. Length grouping reduces padding among dissimilar sequences; global work assignment aims to reduce GPU imbalance.

There are constraints. Training needs the annotated target path and crops in advance, while inference remains sequential when later decisions depend on earlier selections. Scheduling also adds bookkeeping, and the reported method keeps questions as separate causal rows rather than packing them together. Balanced batches improve utilization only insofar as the cost estimate captures actual work. These are efficiency choices around the objective—not a guarantee that every stage or every GPU costs the same.