Focus on Likely Classes for Test-Time Prediction

TL;DR

The paper investigates whether focusing on likely classes for a single in-domain sample can improve predictions. It introduces two test-time fine-tuning methods, iFo and doFo, with a uncertainty-driven gating and a single gradient-step update; iFo tends to amplify shared features among top classes, while doFo can be unstable; experiments across vision and language tasks show iFo yields accuracy gains in a majority of configurations, outperforming several domain-adaptation baselines. The work provides a theoretical motivation via a toy model and contrasts with entropy-minimization; limitations include lack of formal guarantees and variable gains across tasks; practical impact is a lightweight, in-domain TTA that can be selectively activated to improve uncertain predictions.

Abstract

We ask: Can focusing on likely classes of a single, in-domain sample improve model predictions? Prior work argued ``no''. We put forward a novel rationale in favor of ``yes'': Sharedness of features among classes indicates their reliability for a single sample. We aim for an affirmative answer without using hand-engineered augmentations or auxiliary tasks. We propose two novel test-time fine-tuning methods to improve uncertain model predictions. Instead of greedily selecting the most likely class, we introduce an additional step, \emph{focus on the likely classes}, to refine predictions. By applying a single gradient descent step with a large learning rate, we refine predictions when an initial forward pass indicates high uncertainty. The experimental evaluation demonstrates accuracy gains for one of our methods on average, which emphasizes shared features among likely classes. The gains are confirmed across diverse text and image domain models.

Figures (25)

• Figure 1: Overview: Fine-tune based on likely initial classes.
• Figure 2: Our test-time approaches (iFo/doFo): If a prediction exhibits high uncertainty, changes are done to focus on likely (focus) classes either by increasing focus classes or decreasing all others.
• Figure 3: Coefficients for entropy-based minimization are poorly aligned with our rationale.
• Figure 4: Differences in accuracies for $f^{iFo}$ using probability-weighted ($p_c$) or non-weighted likely class terms in the loss function (Eq. \ref{['eq:idoFo']}).
• Figure 5: Difference in Mean Top-k Accuracy $\Delta$Acc$_{k=1,k=2}$ for non-tuned models $f$ across all configurations $O$ grows with lower threshold $d_{1,2}$.