Adaptive Wrong-Token Weighting for Dual Focal Loss in Few-Shot Grid Reasoning

Adaptive Wrong-Token Weighting for Dual Focal Loss in Few-Shot Grid Reasoning Anonymous for Submission Abstract Autoregressive grid-token models used for few-shot adaptation in ARC-style tasks are typically trained using crossentropy loss. However, only a small subset of tokens determines the underlying transformation rule, while most tokens correspond to visually repetitive background structure. This causes gradients on the rule-defining tokens to be diluted when losses are averaged uniformly. Focal and dual focal losses adjust per-token weighting based on confidence, but still aggregate these contributions uniformly across positions, allowing trivial tokens to dominate updates. We introduce a simple aggregation-level modification: we mix dual focal loss over all tokens with dual focal loss over only mispredicted tokens, using a one-cycle cosine schedule to emphasize rule-defining exceptions during the middle of adaptation and reduce emphasis as the model stabilizes. On the ARC-2025 public evaluation benchmark, this approach solves 6/120 tasks compared to 4/120 using cross-entropy and 1/120 using focal or dual focal loss under identical conditions. The results suggest that controlling how token-level losses are aggregated is critical when only a few tokens encode the governing transformation. 1 Introduction The Abstraction and Reasoning Corpus (ARC) evaluates systematic rule inference from only a few input–output examples. In its grid-based formulation, models are trained autoregressively to predict discrete tokens representing colors or structural elements. Most tokens in an ARC grid correspond to background patterns or repeated layout, while the actual transformation rule is often determined by a small number of exception tokens—tokens in which the output deviates from the visually dominant pattern. These exceptions encode the underlying rule, yet they are vastly outnumbered by tokens that merely reflect repetition. Cross-entropy (CE) loss averages per-token negative log-likelihood uniformly and therefore tends to update parameters based primarily on the easy, repetitive structure. Focal loss [2] modifies CE by incorporating a factor (1 −pt)γ, reducing the contribution of well-predicted tokens and increasing emphasis on harder predictions. However, the effect of focal loss depends on the choice of γ, and empirical studies show that overly aggressive reweighting can lead to under-confident predictions, motivating extensions such as Dual Focal Loss [3], which incorporates both the correct logit and strongest competing logit to balance margin and confidence. While Dual Focal Loss improves per-token weighting, it still aggregates losses uniformly across positions. In ARCstyle few-shot adaptation scenarios, this aggregation step remains problematic: because exception tokens are extremely sparse, even appropriately weighted losses can be overwhelmed when averaged with thousands of trivial tokens. Thus, the challenge is not only how to weight individual tokens, but how the weighted losses are combined. We therefore modify the aggregation layer of the loss. Instead of uniform averaging, we mix dual focal loss over all tokens with dual focal loss computed only over mispredicted tokens, using a smooth one-cycle cosine schedule. This allows early training to remain stable, increases emphasis on exception tokens after the model forms an initial hypothesis, and gradually reduces emphasis again near convergence. 1 2 Method For each valid token position i, let pt,i denote the probability assigned to the correct token, and pm,i = max j=ti pj,i denote the strongest competing probability. The dual focal per-token loss is: ℓi= (1 −pt,i +pm,i)γ·(−log pt,i). Let Iall be all valid positions and Iwrong ={i∈ Iall |arg max jpj,i =ti}. We compute: Lall =1 |Iall|X i∈Iall ℓi, Lerr =1 |Iwrong|+ϵX i∈Iwrong ℓi, with ϵ>0to avoid division by zero. A one-cycle cosine schedule controls the mixing: α(t) = 0.5·1−cos(πt) 2, t ∈[0,1]. The final loss is: L(t) = (1 −α(t))Lall +α(t)Lerr. 3 Experimental Setup We evaluate on the ARC-2025 public evaluation set [1]. A pretrained Qwen3-0.6B grid-token model is adapted per task using only the provided demonstrations. Architecture, optimizer, schedules, tokenization, and decoding heuristics are held constant across experiments; only the training loss changes. A task is scored as solved only if the correct output appears among the two allowed final submissions. 4 Results Loss Function Raw Solved Scored Solved (Top-2) Cross-Entropy (CE) 4 2 Focal Loss 1 0 Dual Focal Loss 1 0 Ours (Adaptive Wrong-Token Dual Focal) 6 4 5 Discussion The primary difficulty in ARC-style adaptation arises from the imbalance in how rule-defining tokens contribute to the aggregate gradient. Per-token weighting alone is insufficient when only a handful of positions determine the transformation rule. Adjusting aggregation to temporarily emphasize mispredicted tokens allows the model to override dominant-but-incorrect patterns and converge more reliably. 2 6 Conclusion We introduced an aggregation-level modification to dual focal loss that improves success rates in few-shot grid reasoning. The method is lightweight, architecture-agnostic, and adds no new parameters. The results highlight the importance of shaping how token-level losses are aggregated when only a few tokens encode the correct transformation. References [1] F. Chollet, M. Knoop, G. Kamradt, W. Reade, and A. Howard. ARC Prize 2025. Kaggle, 2025. https:// kaggle.com/competitions/arc-prize-2025. [2] T.-Y. Lin, P. Goyal, R. Girshick, K. He, and P. Doll´ ar. Focal Loss for Dense Object Detection. In ICCV, 2017. [3] L. Tao, M. Dong, and C. Xu. Dual Focal Loss for Calibration. arXiv:2305.13665, 2023. 3