Set Learning for Accurate and Calibrated Models

ArXiv: 2307.02245

🎯 Pitch

Training on sets of examples—where the model guesses which two belong to the same class—dramatically boosts both accuracy and calibration, delivering an 8.59% absolute accuracy gain in 10-shot MNIST over prior best methods, without any post-hoc calibration tricks. The secret is that including odd-class examples alongside the matching pair naturally regularizes logits, slashing overconfidence while standard training on the same data remains miscalibrated.

---

1. Executive Summary

This paper introduces odd-k-out learning (OKO), a training framework that minimizes cross-entropy over sets of examples — each containing a pair of same-class instances and k odd-class instances — rather than over individual samples, and shows on MNIST, FashionMNIST, CIFAR-10, and CIFAR-100 that OKO simultaneously improves both accuracy and calibration over standard training, label smoothing, focal loss, and batch-balancing baselines. The method achieves its strongest gains in limited-data and heavy-tailed class distribution regimes — including an 8.59% absolute improvement over the prior best 10-shot MNIST accuracy — while yielding smoothed logits and lower Expected Calibration Error without any post-hoc temperature scaling or hyperparameter tuning, establishing that training on sets implicitly regularizes against overconfidence only when the set construction includes odd-class examples alongside the pair class.

2. Context and Motivation

The Core Problem: Accuracy and Calibration Are Treated as Opposing Goals

The fundamental tension this paper addresses is that modern deep neural networks achieve high classification accuracy but often produce overconfident, miscalibrated probability estimates — they assign near-certainty to predictions that are frequently wrong. This is not merely a cosmetic issue: in safety-critical domains (medical diagnosis, autonomous driving, scientific experimentation), a model that reports 99% confidence on an incorrect prediction is far more dangerous than one that appropriately hedges with, say, 60% confidence. The paper's central question, stated explicitly in Section 1, is:

"Can we provide a training framework to learn network parameters that simultaneously obtain better accuracy and calibration, especially with class imbalance?"

This framing is important because it rejects the conventional wisdom that accuracy and calibration are orthogonal quantities that must be optimized separately. The prevailing pipeline in practice is: (1) train for accuracy using standard cross-entropy, (2) observe that the resulting model is overconfident, and (3) apply an ad-hoc post-processing fix — typically temperature scaling, which requires a held-out validation set to tune a single scalar parameter that softens the predicted probabilities without affecting the ranking of class predictions (and therefore without changing accuracy). This three-stage pipeline treats calibration as an afterthought rather than a first-class training objective.

Why Miscalibration Happens: The Divergence Problem in Cross-Entropy

The paper identifies a specific mechanism that causes overconfidence under standard training (Section 4 and Appendix C). When minimizing cross-entropy on individual examples with hard labels, the loss function has a structural property: for every class $i$ and every class $j \neq i$, the gradient of the risk $R(F)$ with respect to the logit $F_{i,i}$ (the logit for the true class) is always positive, and with respect to $F_{i,j}$ (off-diagonal logits) is always negative. Concretely:

$\frac{\partial R(F)}{\partial F_{i,i}} > 0 \quad \text{and} \quad \frac{\partial R(F)}{\partial F_{i,j}} < 0 \quad \text{for all } i \neq j$

This means that standard cross-entropy always pushes true-class logits toward $+\infty$ and wrong-class logits toward $-\infty$, regardless of whether the training data provides sufficient evidence to justify such extreme certainty. When a model has enough capacity to memorize the training data — which modern neural networks overwhelmingly do — it will drive logits to extreme values for every training point, producing near-deterministic (0 or 1) output probabilities. These extreme probabilities are poorly calibrated: the model's confidence far exceeds its actual accuracy on unseen data.

This divergence is at the heart of the calibration problem. As the authors note in Appendix C:

"standard cross-entropy loss always encourages the logits of the true class to move towards $\infty$ and the logits of the wrong class towards $-\infty$. As a result, neural networks tend to be overconfident."

Weight decay can partially mitigate this by penalizing large logit magnitudes, but modern architectures typically use little to no weight decay because it harms generalization (Guo et al., 2017). This creates a direct conflict: the mechanism that produces good accuracy (strongly separating logits) is also the mechanism that destroys calibration.

Existing Approaches and Their Shortcomings

The paper situates itself against several existing lines of work on calibration and class imbalance, identifying specific limitations in each.

Post-hoc calibration methods. The most widely used approach is temperature scaling (Platt et al., 1999; Guo et al., 2017): after training, divide all logits by a scalar temperature $\tau > 1$, which softens the softmax output without changing the argmax. This is simple and effective but has three significant drawbacks: (1) it requires a held-out calibration dataset to tune $\tau$, which may not be available in low-data settings; (2) it doesn't change the learned model parameters, meaning the underlying features that drive overconfidence remain unchanged; and (3) it applies a uniform transformation to all predictions regardless of class, which is suboptimal for class-imbalanced settings where minority classes need different treatment than majority classes. More sophisticated post-hoc methods exist — isotonic regression (Zadrozny & Elkan, 2002), Bayesian binning (Naeini et al., 2015), Dirichlet calibration (Kull et al., 2019) — but all share the fundamental limitation of being decoupled from the training process.

Label smoothing. Label smoothing (Müller et al., 2019) addresses miscalibration during training by replacing hard one-hot targets with soft targets: $y_i = 1 - \alpha$ for the true class and $y_j = \alpha / (C-1)$ for the remaining $C-1$ classes, where $\alpha$ is a hyperparameter controlling the smoothing strength. This prevents logits from diverging by ensuring the model never tries to drive the true-class probability all the way to 1.0. However, the paper identifies several limitations. First, it requires careful tuning of $\alpha$ — "fine-tuned parameters for good empirical performance, such as the noise parameter for label smoothing" (Section 1). Second, as shown theoretically in Appendix C (Proposition 3 and surrounding discussion), label smoothing encourages convergence to a fixed minimum that imposes a rigid structure on the logits, potentially limiting model flexibility in regions where high certainty is justified. Third, label smoothing does not address the core issue that calibration is inherently a property of sets of predictions, not individual predictions — it smooths targets but doesn't consider relationships between different examples.

Class imbalance methods. The paper evaluates several approaches for handling class-imbalanced data:

• Error re-weighting (Section B.2): weighting each example's loss inversely proportional to its class frequency, so rare class examples contribute more to the gradient. While principled, this can lead to instability because the weights can become very large for extremely rare classes.
• Focal loss (Lin et al., 2017): modifying cross-entropy to down-weight easy examples and focus training on hard examples. This was designed for object detection and doesn't directly address calibration.
• Batch-balancing (Section B.3): constructing mini-batches by first sampling a class uniformly, then sampling an example from that class, ensuring each class appears equally often on average. Proposition 1 shows this is equivalent in expectation to error re-weighting, but the paper finds it works better in practice — however, it still ignores calibration entirely.

The critical observation is that methods designed for imbalanced accuracy do not improve calibration for minority classes. As the paper notes (Section 1):

"techniques for mitigating the effects of imbalance on classification accuracy do not improve calibration for minority instances and standard calibration procedures tend to systematically underestimate the probabilities for minority class instances"

This is a well-documented phenomenon: when training data for a class is sparse, models tend to be underconfident for that class, even if they correctly classify those examples. Post-hoc temperature scaling applies a uniform transformation that cannot correct for class-specific miscalibration patterns.

The training-calibration disconnect. The deeper problem the paper identifies is conceptual: all prior methods treat calibration as something to fix after or during training via modifications to the loss function or post-processing, but none address the root cause that training on individual examples fundamentally ignores the set-level statistics that define calibration. Calibration is a property of the joint distribution of predictions and outcomes — it's measured by comparing predicted probabilities against empirical frequencies across many predictions. Standard training, which processes examples one at a time, has no mechanism to account for these aggregate statistics. This is the conceptual gap the paper aims to fill.

The Key Insight: Calibration Is a Set-Level Property, So Train on Sets

The paper's core positioning is that calibration should be addressed at the training data level rather than the loss-function or post-processing level. The argument, articulated in Section 1 and the introduction to Section 3, is:

1. Calibration measures aggregate behavior: Expected Calibration Error (ECE) bins predictions by confidence and compares the average accuracy in each bin to the bin's confidence level. This inherently involves sets of examples — you cannot measure calibration on a single test point.
2. Standard training sees no sets: each gradient update depends on a single $(x_i, y_i)$ pair, so the model never receives a signal about how its predictions relate to each other across examples.
3. Therefore: if the training objective itself involves sets of examples, the optimization process will naturally account for the correlations and aggregate statistics that calibration depends on.

This leads to the paper's proposed solution — odd-k-out learning (OKO): present the model with constructed sets containing exactly two examples from the same class (the "pair") and $k$ examples each from different classes (the "odds"), and train the model to identify the pair class by summing logits across the set. The key word in the authors' framing is "naturally" — they argue that set-based training inherently produces calibrated probabilities without needing explicit calibration penalties, temperature tuning, or label smoothing hyperparameters.

Where This Paper Positions Itself in the Literature

The paper draws on several distinct research threads to position OKO as a novel synthesis:

Odd-one-out tasks from cognitive science. The paper explicitly credits the cognitive psychology literature for inspiration (Section 2): the odd-one-out task, where human subjects identify which item in a set doesn't belong or which pair is most similar, has been used for decades to study perceptual similarity (Robilotto & Zaidi, 2004; Hebart et al., 2020). The paper cites prior machine learning work that adapted odd-one-out signals: Fernando et al. (2017) used it for self-supervised video understanding, Locatello et al. (2020) and Mohammadi et al. (2020) used pairwise comparisons as weak supervision, and Muttenthaler et al. (2023b) used human odd-one-out judgments to improve pretrained representations for few-shot learning. However, the paper notes that "none of these works investigated calibration or provided any theory for (odd-one-out) set learning" (Section 2) — this is the gap OKO fills.

Multiple instance learning (MIL). The task of classifying sets of instances is known as multiple instance learning (Carbonneau et al., 2018), where a bag of instances receives a single label. Common approaches include mean pooling (which is "akin to OKO"), max pooling (Feng & Zhou, 2017), attention-based aggregation (Ilse et al., 2018), and permutation-invariant networks (Zaheer et al., 2017). The paper distinguishes OKO from MIL by noting that existing MIL work focuses on predicting the set label and doesn't consider the calibration of the resulting classifier when applied to individual instances at test time. OKO's contribution is showing that set-based training during learning improves single-instance calibration at inference — a finding that is not obvious from the MIL literature.

Theoretical calibration guarantees. The paper engages with the literature on calibration measures and proper scoring rules (Section 4 and Appendix E). It notes that Expected Calibration Error (ECE) — while widely used — has known limitations: it is discontinuous (small changes in the predictor can cause large changes in ECE; Kakade & Foster, 2004; Foster & Hart, 2018), and binning-based approximations introduce additional variance. Recognizing these issues, the paper introduces a novel measure — relative cross-entropy (Definition 1) — that provides per-datapoint calibration information without binning, and proves that it has desirable properties (Lemmas 1–2): it is positive when predictions are overconfident and incorrect, and has expectation zero for perfectly calibrated predictors.

The Practical Motivation: Limited Data and Heavy-Tailed Distributions

A significant portion of the paper's experimental focus is on limited training data regimes and class-imbalanced (heavy-tailed) distributions. This is not arbitrary — it reflects a deliberate targeting of settings where miscalibration is most severe and where post-hoc calibration methods are least applicable. In low-data settings, models overfit to the few available training examples and produce extreme logit values (overconfidence). Post-hoc temperature scaling is problematic here because there may not be enough held-out data to reliably tune the temperature parameter — the calibration set itself suffers from the same scarcity. In class-imbalanced settings, the underrepresentation of minority classes means that calibration is systematically worse for those classes (Wallace & Dahabreh, 2012), and uniform temperature scaling cannot address class-specific miscalibration.

The paper's heavy-tailed experimental setup (Section 5) is deliberately extreme: 90% of probability mass is concentrated in 3 of the 10 (or 100) classes, with the remaining 10% spread across 7 or 97 minority classes. This simulates real-world distributions where a few categories dominate (common diseases, frequent objects) and many categories are rare (rare conditions, specialized items). In these settings, standard cross-entropy training produces models that are simultaneously overconfident on majority classes and underconfident on minority classes — a pattern that OKO is specifically designed to address.

How OKO Reconciles Conflicting Requirements

The paper's positioning can be understood as resolving a tension that prior work accepted as inevitable: that improving calibration requires sacrificing either accuracy (weight decay, which hurts generalization), model flexibility (label smoothing, which imposes a fixed logit structure), or simplicity (temperature scaling, which requires a separate calibration dataset and post-training tuning). OKO claims to circumvent this tradeoff entirely by changing the training data presentation rather than the loss function or post-processing. The mechanism — a set-level objective that implicitly prevents logit divergence — works in the training loop itself, requires no additional hyperparameters, and produces a model that applies to single examples at test time exactly like a standard classifier. This is the paper's key differentiator from prior work and the motivation for both the empirical experiments and the theoretical analyses that follow.

3. Technical Approach

3.1 Reader Orientation

The system is a training procedure that replaces standard single-example cross-entropy with a set-based classification objective during model optimization but leaves the model architecture and inference-time usage completely unchanged. It solves the problem of simultaneously achieving good accuracy and well-calibrated probabilities — especially with limited or imbalanced data — by constructing training "episodes" where each gradient update is computed from a carefully designed set of examples rather than from a single example, which implicitly regularizes logit magnitudes and prevents the divergence toward extreme confidence that standard training encourages.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components:

1. Training Dataset ($\mathcal{D}$) — the original set of $n$ examples, each with an input vector $x_i \in \mathbb{R}^d$ and a class label $y_i \in [C]$. This is organized into per-class partitions $\mathcal{D}_c$ for efficient sampling.

2. OKO Set Sampler (Algorithm 1) — a stochastic procedure that draws from a carefully defined distribution $\mathcal{A}$ over sets $S = ((x'_1, y'_1), \ldots, (x'_{k+2}, y'_{k+2}))$, where exactly two examples share the same class (the "pair class") and $k$ other examples each belong to a distinct, different class (the "odd classes"). This is the core innovation — it defines what the model sees during training.

3. Neural Network $f_\theta$ — a standard classifier that maps inputs to logit vectors in $\mathbb{R}^C$. The architecture (CNN, ResNet) does not change; OKO only modifies how logits are aggregated and how the loss is computed during training. At inference time, the same network is applied to single inputs exactly as in standard training.

4. Set-Level Aggregation — during training only, the logits from all $k+2$ examples in a set are summed element-wise: $f_\theta(S_x) = \sum_{i=1}^{k+2} f_\theta(x'_i)$. This produces a single vector in $\mathbb{R}^C$ that pools evidence across the entire set.

5. OKO Loss Function — either the hard loss $\ell_{\text{hard}}^{\text{oko}}$ that treats the pair class as a one-hot target against the summed logits, or the soft loss $\ell_{\text{soft}}^{\text{oko}}$ that treats the empirical label distribution of all set members as a soft target. The gradient of this loss flows back through the aggregation to update the shared network parameters $\theta$.

Information flow during one training step: The sampler selects a pair class $y'$ uniformly from $[C]$ → it selects two examples from $\mathcal{D}_{y'}$ and $k$ examples from $k$ distinct other classes → the $k+2$ examples are passed through $f_\theta$ to produce $k+2$ logit vectors → these vectors are summed element-wise into a single set-level logit vector → the loss compares this aggregated logit vector against the target (hard label $y'$ or soft label from all $k+2$ set labels) → gradients backpropagate through the sum into the shared network.

3.3 Roadmap for the Deep Dive

• First, the OKO set sampling procedure (Algorithm 1, Section 3), because it defines the data distribution the model sees during training — the set composition (pair + odds) is what drives all downstream effects on calibration and accuracy.
• Second, the OKO loss functions (Equations 1–2), because they define the optimization objective applied to the sampled sets — the choice between hard and soft targets determines how probability mass is distributed in the model's predictions, and understanding why hard targets work better than soft targets requires analyzing both forms.
• Third, the logit aggregation mechanism and its consequences for gradient flow, because the element-wise summation of logits across set members is what distinguishes OKO from standard training — this is the mathematical operation that couples examples together during optimization and produces the smoothing effect.
• Fourth, the theoretical analysis of the OKO loss landscape (Appendix C, Propositions 2–3), because it explains mechanistically why OKO does not force logits to diverge the way standard cross-entropy does — this connects the training procedure to its calibration benefits.
• Fifth, the proof-of-concept analysis on synthetic data (Section 4, Theorem 1), because it provides a minimal, fully tractable example where OKO's behavior can be predicted exactly — showing concretely that OKO assigns higher uncertainty (lower confidence) to low-data regions than standard training would.
• Sixth, the novel calibration measure (Definition 1, Lemmas 1–2, Section 4) introduced alongside OKO, because it provides a per-datapoint measure of excess confidence that complements standard ECE and connects theoretically to the properties OKO optimizes for.

3.4 Detailed, Sentence-Based Technical Breakdown

This is a methods paper with theoretical analysis whose core idea is that training a classifier to predict the majority class in a set containing two same-class examples and $k$ odd-class examples simultaneously improves accuracy and calibration without architectural changes, post-processing, or additional hyperparameters.

---

OKO Set Sampling: Constructing the Training Distribution

The OKO training procedure does not use the original training data directly. Instead, it constructs a new distribution over sets of examples, denoted $\mathcal{A}$, and samples from this distribution to form each training batch. Algorithm 1 in the paper specifies the exact sampling procedure.

Step-by-step construction of one training set $S$:

Step 1: Select the pair class. A class label $y'$ is sampled uniformly from the set of all classes $[C] = \{1, \ldots, C\}$. This class will appear twice in the set — it is called the pair class.

Step 2: Assign the pair labels. The first two elements of the label vector $S_y$ are set to $y'$: formally, $y'_1 \leftarrow y'$ and $y'_2 \leftarrow y'$.

Step 3: Select the odd classes. $k$ additional class labels are sampled uniformly at random without replacement from $[C] \setminus \{y'\}$, i.e., from all classes except the pair class. These are $y'_3, \ldots, y'_{k+2}$. Each is distinct and different from $y'$, ensuring that the set contains exactly $k+1$ distinct classes: one class appears twice, and $k$ classes appear once each.

Step 4: Select examples for each label. For each $i = 1, \ldots, k+2$, a training example $x'_i$ is sampled uniformly at random from $\mathcal{X}_{y'_i}$, the set of all inputs in $\mathcal{D}$ whose label is $y'_i$. Critically, the two examples for the pair class are sampled independently, so $x'_1 \neq x'_2$ — they are two distinct inputs that happen to share the same class.

The output is a tuple $S = (S_x, S_y)$ where $S_x = (x'_1, \ldots, x'_{k+2})$ and $S_y = (y'_1, \ldots, y'_{k+2})$, plus the pair class identifier $y'$ used as the target for the hard loss.

Why this structure matters. An OKO set always contains exactly two examples from the same class and $k$ examples from $k$ different, non-overlapping classes. This composition has several crucial properties:

• The pair class is the majority class in the set — it appears twice, while every odd class appears exactly once. This means the set-level task of "find the most common class" is well-defined and has a unique answer ($y'$) when $k \geq 1$. If $k = 0$, the set would contain only two examples from the same class, and every class in the set would be the majority class — the task degenerates into standard classification on individual augmented examples, losing the set-level structure.
• The construction forces the model to compare examples across classes. To identify the pair class, the model cannot simply memorize per-class prototypes — it must recognize that two specific instances share a class while others do not, which requires representations that capture within-class similarity and between-class differences.
• The selection of examples is uniform within each chosen class, meaning that the original class frequencies do not affect the set composition beyond the initial uniform selection of the pair class. This implicitly balances class representation: each class has probability $1/C$ of being selected as the pair class in any given set, regardless of how many training examples it has. This is the mechanism by which OKO addresses class imbalance during training — rare classes appear as the pair class just as often as common classes.

The constraint $k + 1 \leq C$. The paper notes that $k$ must be chosen so that $k + 1 \leq C$, since we need $k$ distinct odd classes plus one pair class, all different. For datasets with few classes (like binary classification with $C = 2$), $k$ is forced to be at most 1. For datasets with many classes, $k$ can be larger. In practice, the paper sets $k = 1$ for all experiments after finding in ablation studies that "although odd class examples are crucial, OKO is not sensitive to the particular choice of $k$" (Appendix F.4). With $k = 1$, each training set contains three examples: two from the same class and one from a different class.

Computational cost. The paper states that "the training complexity scales linearly in $O(|S|)$ where $|S|$ denotes the number of examples in a set and hence introduces little computational overhead during training." For $k = 1$, each training set has $|S| = 3$ examples, so each forward pass processes three inputs instead of one — a 3× increase in per-step computation. However, since the number of gradient updates is held constant (the paper sets the maximum number of randomly sampled sets to $n_{\text{train}}$ for fair comparison with standard training), the total training time increases only by a small constant factor.

Relationship to batch-balancing. The uniform selection of the pair class in step 1 is analogous to batch-balancing (Algorithm 2 in Appendix B.3), where a class is first selected uniformly, then an example from that class is drawn. However, OKO goes further: the pair class examples are drawn uniformly within the selected class, and the odd class examples provide additional inter-class contrast. This means OKO inherits the class-balancing property of batch-balancing but adds the set-level structure that forces comparison across examples.

---

The OKO Loss Functions: Hard and Soft Targets

After constructing a set $S$, OKO applies a neural network $f_\theta$ to each input in the set, sums the resulting logit vectors, and computes a loss comparing this aggregated logit vector against a target distribution. The paper defines two variants:

Logit aggregation (shared by both losses). For a tuple of input vectors $S_x = (x'_1, \ldots, x'_{k+2})$, the model's output is defined as:

$f_\theta(S_x) := \sum_{i=1}^{k+2} f_\theta(x'_i)$

where $f_\theta(x'_i) \in \mathbb{R}^C$ is the vector of logits for input $x'_i$.

What it computes: the element-wise sum of all $k+2$ individual logit vectors. If the network predicts logit $z_{i,c}$ for class $c$ on input $i$, the aggregated logit for class $c$ is $\sum_i z_{i,c}$. This is equivalent to summing the evidence for each class across all examples in the set.

Why this form: summation is the simplest permutation-invariant aggregation (order of inputs doesn't matter), and it preserves linearity of the logit space — if two inputs both provide evidence for class $c$, their logits add constructively. This means the pair class, which appears twice in the set, naturally receives approximately twice the logit mass as each odd class (assuming the model is well-trained). The summation makes the pair class statistically dominant in the aggregated logits, which is why the model can learn to identify it. Alternative aggregations would be less natural: max pooling would ignore the second pair-class example entirely; average pooling would make the pair class only fractionally dominant; concatenation would break permutation invariance and scale poorly with $k$.

Hard loss. The primary loss used in all reported experiments (the paper found it "always outperforms the soft loss" in ablation studies, Appendix F.5) treats the pair class $y'$ as a one-hot target:

$\ell_{\text{hard}}^{\text{oko}}(S_y, f_\theta(S_x)) := -e_{y'}^\top \log[\text{softmax}(f_\theta(S_x))]$

where $e_{y'} \in \mathbb{R}^C$ is the indicator vector with 1 at index $y'$ and 0 elsewhere, and $\text{softmax}(f_\theta(S_x))_j = \exp(f_\theta(S_x)_j) / \sum_{\ell=1}^C \exp(f_\theta(S_x)_\ell)$.

What it computes: standard cross-entropy between the one-hot encoding of the pair class and the softmax-normalized aggregated logits. The loss is minimized when the softmax probability assigned to class $y'$ is high, i.e., when $\text{softmax}(f_\theta(S_x))_{y'} \approx 1$.

Why this form: this is the standard maximum-likelihood objective for classification, applied at the set level. The model is told "the correct answer for this entire set is $y'$" — it is not told why $y'$ is correct (i.e., that it appears twice while others appear once). The model must learn this statistical regularity from the data distribution, which requires it to count class frequencies within each set. The hard loss provides a single, unambiguous training signal per set.

Soft loss. The alternative loss treats the empirical label distribution of all set members as a soft target:

$\ell_{\text{soft}}^{\text{oko}}(S_y, f_\theta(S_x)) := -\left((k+2)^{-1} \sum_{i=1}^{k+2} e_{y'_i}\right)^\top \log[\text{softmax}(f_\theta(S_x))]$

where $(k+2)^{-1} \sum_i e_{y'_i}$ is a vector in $[0, 1]^C$ whose $c$-th entry is the fraction of examples in $S$ with label $c$. For a set with $k=1$, this is $[2/3, 1/3, 0, \ldots, 0]$ for the pair and odd classes respectively.

What it computes: cross-entropy between a soft target distribution (the empirical class frequencies in the set) and the model's predicted distribution. The model is penalized for not matching these exact proportions in its softmax output.

Why this form exists but is not preferred: the soft loss explicitly tells the model "the pair class should receive $2/(k+2)$ probability, and each odd class should receive $1/(k+2)$," which directly embeds the set's composition into the target. This seems like it should provide more information to the model than the hard loss. However, the paper found empirically that the soft loss produces "predictions that tend to be more uncertain" with "probability mass spread out almost uniformly across classes," leading to worse accuracy (Appendix F.5). The authors hypothesize that the soft loss "transforms the majority class prediction problem into a proportion estimation problem, which may make the model unnecessarily underconfident about the correct class." The hard loss, by contrast, forces the model to commit to a single answer while the set structure implicitly provides the smoothing — the model must learn that the pair class is correct because it appears twice, which is a harder but ultimately more beneficial signal.

Empirical risk. For OKO set sampling $S = (S_x, S_y) \sim \mathcal{A}$, the empirical risk minimized during training is:

$\mathbb{E}_{S \sim \mathcal{A}}\left[\ell_{\text{hard}}^{\text{oko}}(S_y, f_\theta(S_x))\right]$

In practice, this expectation is approximated by sampling sets according to Algorithm 1, forming mini-batches, and applying stochastic gradient descent. The paper notes that it "set[s] the maximum number of randomly sampled sets to the total number of training data points $n_{\text{train}}$ in every setting" to guarantee the same number of gradient updates as standard cross-entropy training for fair comparison.

---

The Loss Landscape: Why OKO Prevents Logit Divergence

The paper provides a theoretical analysis (Appendix C) of how the OKO loss behaves when a model has memorized its training data — i.e., when $f_\theta(x_i) \approx F_{y_i}$, where $F \in \mathbb{R}^{C \times C}$ is a matrix whose row $F_c$ gives the logit vector produced for any input of class $c$.

Standard cross-entropy encourages divergence. For standard cross-entropy, the gradient with respect to any entry of $F$ has a fixed sign:

$\frac{\partial R(F)}{\partial F_{i,i}} > 0 \quad \text{and} \quad \frac{\partial R(F)}{\partial F_{i,j}} < 0 \quad \text{for all } i \neq j$

The gradient always pushes the true-class logit $F_{i,i}$ higher and the wrong-class logits $F_{i,j}$ lower, regardless of how extreme the values already are. This means that without explicit regularization (weight decay, early stopping), gradient descent will drive $F_{i,i} \to +\infty$ and $F_{i,j} \to -\infty$, producing softmax outputs that are one-hot vectors.

The OKO risk functions. The paper derives the population risk for OKO under the assumption $f_\theta(x) = F_y$. Let $\mathcal{Y}_i = \{ Y \mid Y \subset [C], |Y| = k, i \notin Y \}$ be the set of all possible choices of $k$ odd class labels when the pair class is $i$. The soft and hard risks are:

$R_{\text{soft}}(F) = -\sum_{i=1}^C \sum_{Y \in \mathcal{Y}_i} \sum_{j=1}^{k+2} \log\left[\text{softmax}\left(\sum_{\ell=1}^{k+2} F_{Y'_\ell}\right)\right]_{Y'_j}$

$R_{\text{hard}}(F) = -\sum_{i=1}^C \sum_{Y \in \mathcal{Y}_i} \log\left[\text{softmax}\left(\sum_{\ell=1}^{k+2} F_{Y'_\ell}\right)\right]_{i}$

where $Y' = (i, i, Y_1, \ldots, Y_k)$ is the ordered label vector for a set with pair class $i$ and odd classes $Y$.

What these risks compute: for each possible choice of pair class $i$ and odd class set $Y$, the risk evaluates how well the model identifies the correct class from the summed logits. For $R_{\text{hard}}$, only the softmax probability assigned to the pair class $i$ matters. For $R_{\text{soft}}$, the probabilities assigned to every class in the set matter, weighted by their frequency.

Proposition 2: Local convexity and unique minima per logit. The paper proves (Appendix C) that for any fixed pair of class indices $(i, j)$, both $R_{\text{soft}}(F)$ and $R_{\text{hard}}(F)$, when viewed as functions of the single variable $F_{i,j}$ with all other entries held constant, are strictly convex and admit a unique global minimizer.

Why this matters: standard cross-entropy is also convex in each $F_{i,j}$ individually, but its gradient never changes sign, so the minimizer is at $F_{i,i} = +\infty$ and $F_{i,j} = -\infty$ (no finite minimum exists without regularization). OKO's risk, by contrast, has a finite optimal value for each logit when others are held fixed. This is because the logit summation in $f_\theta(S_x)$ couples different entries of $F$ together in the softmax denominator — increasing $F_{i,i}$ affects the loss through multiple summands with different sets $Y$, and the optimal value balances these competing pressures. This acts as an implicit regularizer that prevents logits from diverging.

Proposition 3: Divergence is still possible but gradient descent paths meander. The paper proves that there exist initial values $F^{(0)}$ from which gradient descent on $R_{\text{hard}}$ will cause $F_{i,i}^{(a)} \to +\infty$ and $F_{i,j}^{(a)} \to -\infty$ as $a \to \infty$. So OKO does not guarantee finite logits. However, because the loss is only strictly convex (not monotonic) in each $F_{i,j}$, the gradient path takes a more circuitous route — it "tends to meander rather than directly diverge," as illustrated by a toy example in Figure 6 (Appendix C). This meandering means that under practical optimization (finite steps, non-zero learning rate, stochastic gradients from sampling sets rather than computing the full population risk), the logits reach moderate, well-calibrated values rather than extremes.

Comparison to label smoothing. The paper explicitly distinguishes OKO's effect from label smoothing. Label smoothing replaces hard targets with soft targets $(1-\alpha)e_i + \alpha/C \cdot \mathbf{1}$, which encourages $\exp(F_i)$ to be proportional to this smoothed target. As a result, $F$ is forced toward a unique fixed point (the logits that achieve exactly these softmax proportions). OKO, by Proposition 3, does not force a unique fixed point — "optimizing the (hard) OKO risk still allows $F$ to diverge if that's advantageous." This means OKO provides a more flexible form of regularization: it penalizes extreme logits through the set-level coupling but does not prevent the model from becoming confident when the data genuinely warrants it. The paper summarizes this as:

"the OKO risk strikes a balance between the excessive overconfidence caused by standard cross-entropy and the inflexible calibration of fixed minima in label smoothing."

---

Theorem 1: OKO Is Uncertain in Low-Data Regions

To provide a concrete, fully tractable example of OKO's behavior, the paper analyzes a synthetic binary classification problem with three input values ($x \in \{0, 1, 2\}$) and an imbalanced data distribution defined in Table 1 (Section 4). The distribution has the following structure:

• For $x = 0$: labels are 1 and 2 each with probability $(1-\epsilon)/2$, so most of the data falls here and it is a high-entropy (noisy) region.
• For $x = 1$: label is 1 with probability $\epsilon/2$, label 2 never occurs — this is a low-entropy region (deterministic label given $x$) but with very few samples.
• For $x = 2$: label is 2 with probability $\epsilon/2$, label 1 never occurs — another low-entropy, low-sample region.

Here $0 < \epsilon \ll 1$, representing a tiny fraction of data in the low-entropy regions. A standard cross-entropy model with sufficient capacity would memorize that $x=1 \implies y=1$ and $x=2 \implies y=2$, producing extreme logits for these inputs despite having seen very few examples — a classic case of overconfidence on sparse data.

Theorem 1 (stated formally in Section 4): For all $\epsilon \in (0, 1)$, there exists a minimizer $f_\epsilon$ of the OKO risk (optimizing over all functions from $\{0,1,2\}$ to $\mathbb{R}^2$). Furthermore, as $\epsilon \to 0$, the softmax outputs of these minimizers converge to:

$\text{softmax}(f_\epsilon(0)) \to [1/2, 1/2], \quad \text{softmax}(f_\epsilon(1)) \to [2/3, 1/3], \quad \text{softmax}(f_\epsilon(2)) \to [1/3, 2/3]$

What this says operationally: even though the low-entropy regions ($x=1$ and $x=2$) are perfectly separable and have deterministic labels, OKO does not assign probability 1.0 to the correct class. Instead, it assigns only $2/3$ probability — substantial uncertainty — because these regions contain very little data. The high-entropy region ($x=0$) correctly receives a uniform $[1/2, 1/2]$ prediction.

Why this happens: the OKO objective involves sets that mix examples from different $x$ values. For instance, a set could contain $x=1$ (label 1) as the pair class, with an odd example from $x=2$ (label 2). The model must predict that label 1 appears twice, but it sees that label 2 also appears in the set, and it must calibrate its confidence based on how often such configurations occur. Because $\epsilon$ is tiny, sets involving $x=1$ or $x=2$ are rare in expectation, so the model learns that these inputs provide only weak evidence for their respective classes. The proof (Appendix D) works by expanding the OKO risk into 16 terms (one for each possible configuration of the three inputs in a set with $k=1$), computing the minimizer in closed form, and taking the limit $\epsilon \to 0$.

This theorem directly illustrates the paper's central claim about calibration:

"OKO is still uncertain about these points because they occur infrequently. This may be interpreted as the network manifesting epistemic uncertainty (label uncertainty in an input region due to having few training samples) as aleatoric uncertainty (uncertainty in an input region due to the intrinsic variance in the labels for that region) in the OKO test time outputs."

In standard training, a sample from $x=1$ with label 1 would push $F_{1,1}$ higher and $F_{1,2}$ lower without any counterbalancing signal. In OKO, the same sample appears in sets with other examples, and the optimization balances multiple competing pressures, preventing any one rare input from dominating.

---

Relative Cross-Entropy: A Novel Per-Datapoint Calibration Measure

The paper introduces a new scoring rule to complement ECE and provide per-datapoint insight into calibration errors. This measure is defined in Section 4 and analyzed in Lemmas 1–2.

Definition 1 (Relative Cross-Entropy):

$\text{RC}(P, Q) = H(P, Q) - H(Q)$

where $H(P, Q) = -\sum_{i=1}^C P_i \log(Q_i)$ is the cross-entropy between distributions $P$ and $Q$ (with $P$ typically being a one-hot label $y$ and $Q$ being the predicted distribution $\hat{y}$), and $H(Q) = -\sum_{i=1}^C Q_i \log(Q_i)$ is the entropy of the predicted distribution.

What it computes: the gap between the cross-entropy the model incurs and the entropy of its own prediction. If the model is perfectly confident ($H(Q) \approx 0$), then $\text{RC} \approx H(P, Q)$ — the full cross-entropy. If the model is maximally uncertain ($H(Q) = \log C$), then $\text{RC} = H(P, Q) - \log C$, which is lower. The measure can be negative: if the model is uncertain about a correct prediction, $H(P, Q)$ can be lower than $H(Q)$.

Why this form: the relative cross-entropy isolates excess confidence — the portion of the cross-entropy error that is attributable to the model being more certain than it should be, rather than to the inherent difficulty of the example. Standard cross-entropy conflates these two sources of error: a model that is uncertain about a genuinely ambiguous input incurs high cross-entropy, but this is appropriate and not a sign of miscalibration. RC subtracts the model's own uncertainty $H(Q)$ from the cross-entropy, so that only the mismatch between confidence and correctness remains.

Relationship to KL divergence. The paper notes that RC is "very similar to KL divergence but with a different entropy term." The KL divergence is $D_{\text{KL}}(P \| Q) = H(P, Q) - H(P)$, which subtracts the label entropy $H(P)$. For hard labels, $H(y) = 0$, so KL equals cross-entropy. RC instead subtracts $H(Q)$, making it a function of the model's output rather than the label.

Lemma 1: RC is positive for overconfident errors. For hard labels $y$, if the predicted probability assigned to the true class satisfies $\hat{y}_y \leq 1/|C|$ (i.e., the prediction is worse than random guessing for the correct class), then $\text{RC}(y, \hat{y}) \geq 0$.

What this means: when a model makes an incorrect prediction with high confidence (the dangerous case), RC is positive and large. When a model makes an incorrect prediction with appropriately low confidence, RC can be negative — the model is appropriately uncertain about its wrong answer.

Lemma 2: Zero expectation for perfectly calibrated predictors. If $\hat{y}$ is a predictor that is perfectly calibrated across the data distribution $\mathcal{D}$, meaning $\mathbb{E}[y \mid \hat{y}(x) = q] = q$ for all $q$, then the average relative cross-entropy is $\mathbb{E}_{(x,y) \sim \mathcal{D}}[\text{RC}(y, \hat{y}(x))] = 0$.

What this means: perfect calibration is characterized by RC having zero mean across the dataset. Positive mean RC indicates systematic overconfidence; negative mean RC indicates systematic underconfidence. This provides a simple scalar diagnostic: if the average RC is positive, the model tends to be overconfident; if negative, underconfident.

What makes this useful: unlike ECE, which requires binning predictions by confidence and comparing against empirical accuracy in each bin, RC is computed per datapoint and averaged, providing a smooth, differentiable measure that can be tracked during training. It also provides per-datapoint information: a scatter plot of $\text{RC}(y_i, \hat{y}_i)$ against $H(\hat{y}_i)$ reveals which predictions are driving miscalibration (e.g., high confidence + high RC = overconfident error).

The paper uses RC empirically to compare OKO against baselines (Section 5, Figure 5, Table 3), measuring the mean absolute difference between the average cross-entropy $\bar{H}(P, Q)$ and the average entropy $\bar{H}(Q)$ across the test set. A low MAE indicates good calibration — the model's cross-entropy error closely matches its own uncertainty. OKO achieves the lowest MAE in most settings, confirming that its predictions exhibit less excess confidence than competing methods.

---

Training Procedure at Inference Time

A key practical property of OKO is that it requires zero changes at inference time. The paper emphasizes this repeatedly (Section 1, Section 3, Appendix A):

• During training: inputs are grouped into sets of size $k+2$, logits are summed across the set, and the loss compares the aggregated logit vector against the target.
• During inference: the model is applied to single examples exactly like any standard classifier. The logit summation is not used; the network simply computes $f_\theta(x)$ and outputs $\text{softmax}(f_\theta(x))$ as its prediction.

The paper provides JAX code in Appendix A showing how this is implemented: a single classification head (a dense layer mapping from representation dimension to $C$ classes) is used with a boolean train flag — if True, inputs are reshaped into sets and logits are summed across the set dimension; if False, the classification head is applied directly to individual inputs. This means OKO-trained models can be deployed identically to standard models with no additional computational overhead at inference.

---

Design Choices and Their Justifications

$k=1$ as the default. The paper ablates $k \in \{0, 1, 2, 3\}$ (Appendix F.4) and finds that:

• $k=0$ (no odd class — just a pair of same-class examples) performs substantially worse, demonstrating that "odd class examples are crucial" for both accuracy and calibration.
• Larger $k$ values perform similarly to $k=1$, suggesting that one odd class provides sufficient inter-class contrast. $k=1$ is computationally cheapest (3 examples per set), so it is used for all main experiments.

Hard loss over soft loss. The soft loss explicitly encodes the set composition in the target distribution, which intuitively seems more informative. However, the paper found it produces overly uncertain predictions (probability mass spread out nearly uniformly across classes) and worse accuracy. The interpretation is that the hard loss forces the model to discover the set structure (the pair class appears twice) rather than having it spoon-fed through the target, which leads to better representations and more appropriate confidence calibration.

Auxiliary odd-class prediction head. The paper notes that "preliminary experiments have shown that generalization performance can be boosted by predicting the odd class using an additional classification head" (Section 5). For $k=1$, the model has a second classification head that predicts which class is the odd class (in addition to the main head that predicts the pair class). This auxiliary head is discarded at inference time — it serves only as an additional training signal that encourages the model to distinguish all classes in the set, not just the majority. This is a standard multi-task learning technique adapted to the OKO framework.

Matching the number of gradient updates. To ensure fair comparison with standard training, the paper generates exactly $n_{\text{train}}$ OKO sets per epoch, where $n_{\text{train}}$ is the number of original training examples. Since each OKO set contains $k+2$ examples, the effective number of examples seen per epoch is $(k+2) \times n_{\text{train}}$, but the number of gradient updates is the same. This controls for optimization budget: any improvement from OKO is due to the set-level objective, not to seeing more data or taking more optimization steps.

No additional hyperparameters. Unlike label smoothing (which requires tuning $\alpha$) or temperature scaling (which requires tuning $\tau$ on a held-out set), OKO introduces no new hyperparameters beyond $k$, which is fixed at 1 for all experiments after initial ablation. The only design choices are the loss variant (hard vs. soft) and whether to include the auxiliary odd-class head — both are binary choices, not continuous parameters requiring grid search. This makes OKO practical for low-data settings where held-out calibration data may be insufficient for hyperparameter tuning.

4. Key Insights and Innovations

Innovation 1: Calibration as a Set-Level Property That Should Be Optimized at the Set Level During Training

The paper's most fundamental conceptual move is redefining calibration from a post-hoc measurement or loss-function modification into a training data structure problem. This reframing is what makes OKO more than just another regularization technique.

Before this work, the field operated under an implicit assumption that calibration must be addressed either by modifying what the model optimizes (label smoothing, focal loss, weight decay) or by transforming outputs after training (temperature scaling, isotonic regression). Both approaches accept the basic paradigm of single-example empirical risk minimization and try to compensate for its calibration-destroying tendencies after the fact. Label smoothing softens the targets; temperature scaling softens the outputs — but neither addresses the root cause, which is that training on individual examples provides no signal about aggregate prediction statistics.

The paper's diagnostic insight is deceptively simple: calibration is defined by comparing predicted probabilities against empirical frequencies across many predictions — it is inherently a property of sets. ECE, the field's standard calibration metric, works by binning predictions by confidence level and measuring the gap between average confidence and average accuracy in each bin. These averages are meaningless for single examples. Yet standard training processes examples one at a time, computing gradients from isolated $(x_i, y_i)$ pairs. The optimization never "sees" relationships between predictions on different examples, so it has no mechanism to learn that being 99% confident across 100 predictions should translate to roughly 99 of them being correct.

OKO closes this gap directly: it constructs training episodes where the model must reason about sets of examples to solve the task, and the gradient signal from each episode depends on the aggregate statistics of that set. When the model processes a set with two same-class examples and one odd-class example, the loss compares the summed logits against the pair class label. To minimize this loss, the model must produce logits that, when summed, make the pair class dominant. This couples the optimization of individual example representations through the set-level aggregation, creating an implicit pressure for the model's confidence to reflect how often a class actually appears in training sets — which, by the design of the sampling procedure, reflects properties of the underlying data distribution like class frequency and within-class coherence.

The significance of this reframing extends beyond OKO's specific implementation. It suggests that any training objective that operates on sets rather than individuals can induce calibration-like regularization without explicit calibration penalties. This opens a new design axis for future work: rather than asking "what loss function or post-processing produces calibration?", researchers can ask "what set construction induces the right calibration properties for a given task?" The paper's finding that $k=0$ (sets with only a pair, no odd class) performs substantially worse than $k \geq 1$ (Appendix F.4) demonstrates that the specific set composition matters — the presence of contrastive examples from different classes is what drives the calibration benefit, not merely seeing multiple examples. This is a more nuanced insight than "train on sets" and points toward a principled theory of set design for calibration that the paper initiates but does not complete.

Innovation 2: Implicit Entropic Regularization Through Set-Level Coupling, Without Hyperparameters or Fixed Targets

The paper's second distinctive contribution is demonstrating — both theoretically (Appendix C) and empirically — that training on appropriately constructed sets produces a form of regularization that is qualitatively different from both label smoothing and weight decay, and that this regularization emerges from the optimization dynamics rather than from an explicit penalty term.

Label smoothing prevents logit divergence by replacing hard one-hot targets with soft targets: the model is told "aim for probability $1-\alpha$ on the correct class, not 1.0." This works but imposes a rigid structure — the logits are driven toward a fixed point where $\text{softmax}(F_i) \approx (1-\alpha)e_i + \alpha/C \cdot \mathbf{1}$. Proposition 3 in Appendix C proves that OKO's risk function does not have this property: gradient descent on the hard OKO loss can still cause logits to diverge to $\pm \infty$ from certain initializations. This might seem like a weakness, but the paper frames it as a strength — it means OKO provides flexible regularization that penalizes extreme logits in low-data regions (Theorem 1 shows predictions converge to $[2/3, 1/3]$ rather than $[1, 0]$ on sparse but separable data) while preserving the ability to become highly confident when the data genuinely warrants it.

The mechanism that produces this flexibility is the coupling of logit entries through set-level summation. In standard cross-entropy, the gradient with respect to $F_{i,i}$ depends only on whether example $i$ is correctly classified — it's always positive, pushing the logit higher. In OKO, the gradient with respect to $F_{i,i}$ depends on all examples in the set through the softmax denominator of the summed logits. When class $i$ appears as the pair class in a set that also contains an odd-class example of class $j$, increasing $F_{i,i}$ increases the denominator of the softmax for class $j$'s contribution, which can increase the loss from sets where $j$ is the pair class. This creates competing pressures that find an equilibrium at moderate logit values — but only when the data distribution provides sufficient counterbalancing signals. In data-rich regions where class $i$ appears as the pair class far more often than it appears as an odd class, the pressure to increase $F_{i,i}$ dominates, and the model can become confident.

This is fundamentally different from how weight decay works (which uniformly penalizes all large logit magnitudes regardless of class or data frequency) and from how label smoothing works (which imposes the same soft target structure everywhere). The paper captures this distinction precisely: "the OKO risk strikes a balance between the excessive overconfidence caused by standard cross-entropy and the inflexible calibration of fixed minima in label smoothing." The regularization is data-dependent and adaptive — it's strongest where data is sparse (Theorem 1) and weakest where data is abundant, exactly matching the epistemic uncertainty patterns that a well-calibrated model should exhibit.

The practical consequence — that OKO requires no hyperparameter tuning while label smoothing requires careful selection of $\alpha$ — is a direct result of this implicit mechanism. Label smoothing's $\alpha$ is a global knob that controls the softness of all targets uniformly; getting it right requires held-out data and grid search. OKO's regularization strength is determined by the data distribution itself, through the relative frequencies with which different class combinations appear in sampled sets. This makes OKO self-tuning in a sense that prior calibration methods are not.

Innovation 3: Epistemic Uncertainty Manifested as Aleatoric Uncertainty — A New Diagnostic for Sparse-Data Overconfidence

The paper's third insight is a specific diagnostic finding about how OKO expresses uncertainty, which emerges from Theorem 1 and is validated empirically across the entropy distribution analyses (Figures 4, 8, 9). The phenomenon is that OKO converts epistemic uncertainty (uncertainty due to limited data in a region) into aleatoric uncertainty (uncertainty expressed in the predicted probability distribution for that region) at test time.

In standard training, a model that sees only a few examples of some input region with consistent labels will memorize those examples and become overconfident — it assigns probability near 1.0 to the observed label, expressing near-zero aleatoric uncertainty, even though the epistemic uncertainty (the model's uncertainty about what the true label distribution is, given limited data) should be high. This is the classic overconfidence problem that makes neural networks dangerous in low-data settings: they don't know what they don't know.

Theorem 1 demonstrates mathematically that OKO inverts this pattern. On the synthetic problem where $x=1$ always has label 1 but appears with probability $\epsilon/2 \ll 1$, the OKO-optimal predictor assigns probability only $2/3$ to the correct label — it expresses substantial aleatoric uncertainty (the predicted distribution is not concentrated on a single class) even though the true label distribution in that region has zero entropy. The paper explicitly names this phenomenon:

"This may be interpreted as the network manifesting epistemic uncertainty (label uncertainty in an input region due to having few training samples) as aleatoric uncertainty (uncertainty in an input region due to the intrinsic variance in the labels for that region) in the OKO test time outputs."

The significance of this finding is that it provides a principled diagnostic for distinguishing well-calibrated from overconfident models in sparse-data regions. A standard model's entropy distribution for correct predictions will be concentrated near 0 (high confidence) regardless of how much training data exists for those classes; an OKO model's entropy distribution will shift toward $\log C$ (high uncertainty) for classes or regions with sparse data. This is visible in the empirical entropy plots (Figure 4 for FashionMNIST heavy-tailed, Figures 8–9 for CIFAR variants): OKO shows more entropy mass in the middle range for correct predictions compared to vanilla training, indicating appropriate caution on less-represented inputs.

The paper is careful to note that "the desirability of this property may be somewhat dependent on the application. While it is conceivable that such regions do indeed have low aleatoric uncertainty, and that the few samples do indeed characterize the entire region, for safety-critical applications it is often desirable for the network to err towards uncertainty." This is an honest acknowledgment that converting epistemic to aleatoric uncertainty is a choice with tradeoffs — in applications where the sparse data truly does represent the full distribution (e.g., a rare but perfectly identified medical condition), the induced uncertainty may be overly conservative. But in the more common case where sparse data signals incomplete coverage, this property is exactly what safe deployment requires.

Innovation 4: Relative Cross-Entropy as a Sample-Level, Expectation-Zero Diagnostic for Excess Confidence

The paper introduces a new calibration scoring rule — relative cross-entropy (RC) — that, while not as prominent as OKO itself, represents a genuine methodological contribution to the calibration measurement literature. Its distinctiveness lies in providing per-datapoint, non-negative (for overconfident errors) calibration information with zero expectation under perfect calibration, without requiring binning or population-level aggregation.

The standard metric, ECE, has well-documented flaws that the paper acknowledges (Appendix E): it is discontinuous (small predictor changes can cause large ECE jumps), requires arbitrary binning choices that affect the result, and provides only a population-level summary without identifying which predictions are miscalibrated. These limitations are not merely aesthetic — they mean ECE cannot be used as a training signal (it's non-differentiable and discontinuous) and cannot guide per-example diagnostics in deployment (you can't say "this specific prediction is overconfident" using ECE alone).

RC addresses these by subtracting the model's own entropy from its cross-entropy: $\text{RC}(y, \hat{y}) = H(y, \hat{y}) - H(\hat{y})$. Lemma 1 proves that when a prediction is overconfident and incorrect ($\hat{y}_y \leq 1/C$), RC is non-negative — it detects the dangerous case where the model is simultaneously wrong and certain. Lemma 2 provides the theoretical anchor: for any perfectly calibrated predictor, $\mathbb{E}[\text{RC}] = 0$, giving a clear null hypothesis. A positive mean RC signals systematic overconfidence; a negative mean signals underconfidence.

What makes this more than just another metric is that it operates at the individual prediction level while retaining population-level interpretability. You can plot RC against $H(\hat{y})$ for each test point (as the paper does implicitly through Figure 5 and Table 3) and immediately see which predictions are driving miscalibration — those with low entropy (high confidence) and positive RC are overconfident errors; those with high entropy and negative RC are appropriately uncertain about wrong answers. This per-datapoint diagnostic capability is not available from ECE or Brier score.

The empirical results using RC (Figure 5, Table 3) validate its utility: OKO achieves the lowest mean absolute difference between average cross-entropy and average entropy across almost all settings, confirming that its predictions exhibit the least excess confidence. This is consistent with the ECE trends (Figure 1B, Figure 10) but provides a complementary signal — while ECE measures miscalibration of confidence bins, RC measures the alignment between overall uncertainty and incurred error, capturing a slightly different aspect of calibration quality. The availability of both measures in the paper strengthens the overall calibration claims and provides a template for future work to evaluate calibration at multiple granularities.

The RC measure represents an incremental but practically meaningful advance: it doesn't replace ECE but complements it, filling a gap — per-datapoint calibration diagnostics with clear theoretical properties — that the field had not previously addressed. Its connection to the OKO framework is natural (OKO's set-level objective implicitly optimizes for low excess confidence) but RC is independently useful as an evaluation tool for any calibration method.

5. Experimental Analysis

Evaluation Methodology

Dataset. All experiments use four standard image classification benchmarks: MNIST (10 classes, grayscale digits), FashionMNIST (10 classes, grayscale clothing items), CIFAR-10 (10 classes, color images), and CIFAR-100 (100 classes, color images). For each dataset, the paper uses the official test set for evaluation (held constant across all training configurations) and varies the number of training data points to study low-data regimes. The specific training set sizes are listed in Figure 2's x-axis labels: for MNIST, data points range from 100 to 5,000; for FashionMNIST, 200 to 10,000; for CIFAR-10, 400 to 20,000; and for CIFAR-100, 5,000 to 30,000 (with "Uniform" representing the full balanced training set). For heavy-tailed experiments, the same total counts are used, but the class distribution is modified so that 90% of probability mass is concentrated uniformly across 3 overrepresented classes, with the remaining 10% spread uniformly across the remaining 7 or 97 underrepresented classes.

Base model(s). The paper uses a simple randomly-initialized CNN for MNIST and FashionMNIST, and ResNet18 and ResNet34 architectures (He et al., 2016) for CIFAR-10 and CIFAR-100 respectively. These are standard, well-understood architectures chosen to demonstrate that OKO's benefits are not architecture-specific. The authors do not use pretrained weights — all models are trained from scratch, which is important for evaluating generalization in low-data regimes where pretraining would confound the comparison.

Metrics. Three categories of metrics are reported:

• Classification accuracy: test set accuracy (%), computed as the fraction of test examples where the model's argmax prediction matches the ground-truth label. Reported per training configuration and averaged across five random seeds.
• Expected Calibration Error (ECE): the standard binned calibration metric, computed using the one-vs-all extension for multi-class classification (Zadrozny & Elkan, 2002). Predictions are binned by confidence, and ECE measures the weighted average gap between accuracy and confidence in each bin. Lower ECE indicates better calibration.
• Relative cross-entropy (RC) analysis: the paper's novel metric (Definition 1), quantified as the mean absolute difference (MAE) between the average cross-entropy $\bar{H}(P, Q)$ and the average entropy $\bar{H}(Q)$ across the test set. This captures how closely the model's incurred error tracks its own uncertainty — a low MAE indicates that the model is not systematically overconfident or underconfident. Reported in Table 3 and visualized in Figure 5.

Baselines. Seven comparison methods are evaluated (Section 5, "Training methods"):

1. Vanilla: standard maximum-likelihood estimation with cross-entropy loss on individual examples (Equation 4 in Appendix B.1).
2. Vanilla + LS: vanilla training with label smoothing (Müller et al., 2019), using smoothing parameter $\alpha = 0.1$ (shown in reliability diagram captions and Figure 1 legend).
3. Weighted CE: cross-entropy with per-example weights inversely proportional to class frequency (Equation 5 in Appendix B.2), a standard re-weighting approach for class imbalance (Ting, 2000; Khan et al., 2018).
4. Focal Loss: the focal loss from Lin et al. (2017), designed to down-weight easy examples and focus on hard ones — originally proposed for object detection, evaluated here as a general calibration/imbalance method.
5. Batch-balancing (BB): mini-batch construction where classes are first sampled uniformly, then an example is drawn from the selected class (Algorithm 2 in Appendix B.3). This ensures each class appears equally often in expectation, regardless of its training set frequency.
6. BB + LS: batch-balancing combined with label smoothing ($\alpha = 0.1$).
7. BB + TS: batch-balancing with post-hoc temperature scaling ($\tau = 2.0$), applied after training using a held-out validation set. This is included only in calibration comparisons, not accuracy, since temperature scaling does not change the argmax prediction.

For generalization performance analyses, temperature scaling is excluded because it yields identical accuracy to BB alone. The paper also includes majority voting implicitly through the batch-balancing variants, but does not report a separate "majority voting" baseline.

Generation budget / compute accounting. The paper controls for optimization budget by matching the number of gradient updates across methods. For OKO, the maximum number of randomly sampled sets per epoch is set to $n_{\text{train}}$, the total number of training data points. Since each OKO set with $k=1$ contains 3 examples, the effective number of examples processed is $3 \times n_{\text{train}}$, but the number of forward/backward passes equals the number of gradient updates in standard training. This means OKO's per-epoch computation is approximately 3× higher than vanilla training, but the optimization budget (number of parameter updates) is held constant. The paper argues this is a conservative comparison — OKO's gains come from the set-level objective, not from taking more optimization steps.

Cross-validation / statistical protocol. All experiments are run across five random seeds for every number of training data points and every method. Error bands in figures (e.g., Figure 2, Figure 10) depict 95% confidence intervals computed over these five seeds. Results are reported as averages across seeds. The test set is held fixed across all configurations to ensure comparability. For heavy-tailed experiments, the class distribution is imposed during training, but evaluation is on the original (uniform) test distribution — the model must generalize from an imbalanced training set to a balanced test set, which is the standard and practically relevant evaluation protocol for class imbalance methods. For the RC analysis (Figure 5, Table 3), each point represents a specific seed at a specific training data count, and the MAE is computed across these configurations.

---

Main Quantitative Results

Generalization Performance: OKO Improves Accuracy Across the Board, Especially with Class Imbalance and Limited Data

The paper's headline accuracy results appear in Figure 2 and are summarized numerically in Table 2. The central finding is that OKO achieves the highest average test accuracy across all four datasets and both class distribution settings (uniform and heavy-tailed), with improvements being most pronounced in exactly the challenging regimes that motivate the work.

Aggregate performance (Table 2). Averaged across all training data point settings shown in Figure 2, OKO outperforms every baseline on every dataset-distribution combination:

The pattern is consistent: OKO is never worse than the best baseline, and the margins are largest on MNIST (where absolute accuracy is high enough to see meaningful gaps) and on heavy-tailed distributions (where class imbalance makes standard training particularly brittle). The CIFAR-100 heavy-tailed setting is notable for its low absolute numbers — all methods struggle, but OKO still edges out BB + LS by 0.17 percentage points.

Low-data regimes (Figure 2, top row — uniform distributions). At the smallest training set sizes, OKO's advantage is stark:

• 10-shot MNIST (100 total training points): OKO achieves an average test accuracy of 87.62%, with the best random seed reaching 90.14%. This improves upon the previously reported best 10-shot MNIST accuracy by 8.59 percentage points (Liu et al., 2022). For context, the next best method in the paper's comparison — vanilla training — achieves roughly 82–83% at the same data count (read from Figure 2).
• 20-shot MNIST (200 training points): OKO improves upon Liu et al. (2022)'s best by 2.85 percentage points.
• 50-shot MNIST (500 training points): OKO improves upon Liu et al. (2022)'s best by 1.81 percentage points.
• As training data increases, all methods converge toward similar performance, but OKO maintains a consistent lead: at 5000 training points on MNIST, OKO reaches approximately 97.5% vs. roughly 96.5% for batch-balancing and vanilla training.

Class-imbalanced regimes (Figure 2, bottom row). The heavy-tailed setting reveals OKO's strongest advantages:

• On FashionMNIST heavy-tailed, OKO (red line) tracks above all other methods across the entire data range, reaching approximately 89% at 10,000 training points vs. approximately 85% for the next best method (batch-balancing + LS).
• On CIFAR-10 heavy-tailed, OKO and BB are essentially tied, with both reaching approximately 67% at 20,000 training points. However, OKO shows a slight edge at intermediate data counts (1000–5000 range).
• On CIFAR-100 heavy-tailed, all methods perform poorly (13–14% at 30,000 training points), but OKO achieves the highest average. This is the hardest setting — 100 classes, only 3 of which are well-represented — and the absolute performance is low enough that differences between methods are small in magnitude.
• On MNIST heavy-tailed, OKO dramatically outperforms all baselines: at 500 training points, OKO achieves approximately 88% vs. approximately 81% for batch-balancing (the next best), a 7 percentage point gap. Even at 5000 points, OKO maintains a lead of roughly 2 percentage points.

What drives OKO's accuracy advantage? The paper does not provide a decomposition of accuracy gains by class (majority vs. minority), so we cannot determine from the reported results whether OKO improves both majority and minority class accuracy, or primarily prevents catastrophic failure on minority classes while maintaining majority-class performance. However, the heavy-tailed results — where minority classes make up a larger fraction of the test set due to balanced evaluation — suggest that OKO's gains come substantially from better minority-class recognition. This is consistent with the mechanism described in Section 3: by uniformly sampling the pair class, OKO presents minority classes as the training target just as often as majority classes, preventing the model from ignoring them.

---

Calibration Performance: OKO Achieves Lower ECE Than All Baselines in Most Settings

The calibration results are presented through three complementary visualizations and a summary metric.

Reliability diagrams (Figure 3 for uniform, Figure 7 for heavy-tailed). These diagrams plot accuracy vs. confidence, binned across predictions, with perfect calibration corresponding to the diagonal. The diagrams are averaged across all training data point settings and all five random seeds for each method. The key observations:

• Vanilla training consistently lies below the diagonal across all datasets, especially at high confidence levels — the classic overconfidence pattern. On MNIST uniform (Figure 3), vanilla training shows confidence bins at 0.9–1.0 where accuracy is only 0.85–0.90.
• Label smoothing pulls the curve closer to the diagonal but often overcorrects, producing underconfidence at intermediate confidence levels — the curve lies above the diagonal in the 0.5–0.8 confidence range for FashionMNIST and CIFAR-10.
• Batch-balancing combinations (BB + LS, BB + TS) generally track the diagonal well, producing the most competitive calibration among baselines.
• OKO lies closest to the diagonal across all four datasets, with less systematic deviation than any baseline. On MNIST uniform, OKO's curve is essentially on the diagonal from 0.3 to 1.0 confidence; on FashionMNIST, it shows only slight underconfidence at low confidence levels but tracks perfectly above 0.5.

The heavy-tailed reliability diagrams (Figure 7 in Appendix F.1) show similar patterns but with larger gaps between methods. On MNIST heavy-tailed, only OKO and BB + TS track the diagonal; vanilla training is far below it (dramatic overconfidence), while label smoothing is far above it at low confidence (underconfidence). On CIFAR-100 heavy-tailed, all methods struggle, but OKO's curve is visibly closer to the diagonal than any baseline.

Entropy distributions (Figure 4 for FashionMNIST heavy-tailed; Figures 8–9 in Appendix F.2 for other settings). These plots show the distribution of predictive entropies $H(\hat{y})$ partitioned into correct and incorrect predictions, further split by whether the example came from a "mode" (overrepresented) or "tail" (underrepresented) class during training. The key patterns:

• Vanilla training on FashionMNIST heavy-tailed (Figure 4): correct predictions from mode classes show entropy concentrated near 0 (high confidence, appropriate), but correct predictions from tail classes also show entropy near 0 — the model is overconfident about tail class predictions that happen to be correct. Even more problematically, incorrect predictions from mode classes show significant mass near 0 entropy — overconfident errors.
• Label smoothing shifts entropy mass rightward (toward higher uncertainty) for both correct and incorrect predictions, but this is a blunt instrument — it makes the model underconfident about correct mode-class predictions.
• OKO: correct predictions from mode classes show entropy mass concentrated near 0 (appropriate confidence), while correct predictions from tail classes show visibly more entropy mass in the middle range (0.5–1.0) — the model is appropriately more cautious about rare classes. Incorrect predictions, regardless of class, show entropy mass concentrated near $\log C$ (high uncertainty) — the model knows when it's likely wrong. This is the "epistemic uncertainty manifested as aleatoric uncertainty" pattern that Theorem 1 predicts.

ECE as a function of training data (Figure 1B in Section 1; Figure 10 in Appendix F.3). The paper's central calibration figure is reproduced at larger scale in Figure 10, which plots ECE vs. number of training data points for each method. Key findings:

• Uniform distributions (Figure 10, top row): On MNIST, OKO achieves the lowest ECE across almost all data counts, with values around 0.02–0.04 at 5000 points vs. 0.06–0.10 for vanilla and 0.04–0.06 for label smoothing. On FashionMNIST, OKO and label smoothing are comparable (both around 0.04–0.06 at 10,000 points). On CIFAR-10, OKO's ECE (approximately 0.08–0.10 at 20,000 points) is slightly higher than BB + LS (0.05–0.07). On CIFAR-100, BB + LS achieves the lowest ECE (0.10–0.12) while OKO reaches 0.14–0.16 — still better than vanilla (0.20–0.25) but not the best.
• Heavy-tailed distributions (Figure 10, bottom row): On MNIST and FashionMNIST, OKO achieves the lowest ECE by clear margins. On MNIST heavy-tailed at 500 training points, OKO's ECE is approximately 0.08 vs. 0.25+ for vanilla and 0.15 for BB + LS. On CIFAR-10 heavy-tailed, OKO (0.15–0.20) is comparable to BB + LS. On CIFAR-100 heavy-tailed, all methods show high ECE (0.30–0.45), with BB + LS slightly lower than OKO at some data counts.

The ECE-vs-error tradeoff (Figure 1B, top and bottom panels). Figure 1B in the paper's introduction plots ECE against test classification error, with each point representing one seed at one training data count, and dashed lines showing linear regression fits. This visualization addresses a crucial question: does OKO achieve lower ECE simply because it achieves higher accuracy (which mechanically reduces ECE since well-classified examples tend to have higher confidence)? The answer is yes — part of OKO's calibration improvement is coupled to its accuracy improvement — but the regression lines show that OKO's ECE is lower than expected from its error rate alone. OKO's regression line (red) lies below all baseline lines, meaning that for a given classification error, OKO achieves a lower ECE. This is true for both uniform and heavy-tailed distributions across all four datasets, with the gap being largest on MNIST and FashionMNIST.

---

Relative Cross-Entropy Analysis: OKO Achieves the Lowest Excess Confidence

RC scatter plots (Figure 5, Table 3). Figure 5 plots the average cross-entropy $\bar{H}(P, Q)$ against the average entropy $\bar{H}(Q)$ for different numbers of training data points, with each point representing one seed at one data count. A perfectly calibrated model would have $\bar{H}(P, Q) = \bar{H}(Q)$ — the points would lie on the diagonal line $y = x$. Deviation above the diagonal indicates overconfidence (cross-entropy exceeds entropy); deviation below indicates underconfidence.

• MNIST uniform (top-left): OKO's points (red) cluster tightly around the diagonal. Vanilla training points lie substantially above the diagonal (overconfidence). Label smoothing points lie below the diagonal at low entropy values (underconfidence).
• FashionMNIST uniform (top-second): Similar pattern — OKO tracks the diagonal; vanilla is above it (overconfident); label smoothing is below at low entropy.
• CIFAR-10 uniform (top-third): OKO again closest to diagonal. Batch-balancing with LS also tracks well.
• CIFAR-100 uniform (top-right): OKO and BB + LS are comparable, both tracking near the diagonal with some scatter.
• Heavy-tailed settings (bottom row): The scatter increases for all methods. On MNIST heavy-tailed, OKO remains close to the diagonal while vanilla training points drift far above it (extreme overconfidence at low entropy). On CIFAR-100 heavy-tailed, all methods show substantial scatter, but OKO's points are generally closer to the diagonal than vanilla or weighted CE.

Quantitative MAE summary (Table 3). The mean absolute difference between $\bar{H}(P, Q)$ and $\bar{H}(Q)$, averaged over the test set and all training data counts, confirms the visual patterns:

OKO achieves the lowest MAE in 5 of 8 settings. In the settings where it does not win (FashionMNIST heavy-tailed, CIFAR-10 heavy-tailed, CIFAR-100 heavy-tailed), the best baseline is either Focal Loss, BB + LS, or Weighted CE — methods that explicitly target class imbalance through loss modification. However, in all cases, OKO's MAE is substantially lower than vanilla training (e.g., FashionMNIST heavy-tailed: OKO 0.334 vs. vanilla 1.075; CIFAR-100 heavy-tailed: OKO 1.164 vs. vanilla 2.638), confirming that OKO substantially reduces excess confidence even when it is not the absolute best.

What RC adds beyond ECE. The RC analysis reveals a dimension of calibration that ECE does not capture: the alignment between a model's overall uncertainty and its incurred error. ECE measures whether confidence bins are accurate on average, but a model can have low ECE while being systematically wrong about which predictions are uncertain (e.g., being overconfident on some examples and underconfident on others in ways that cancel out in bin averages). RC is a per-datapoint measure that cannot cancel out — if a model is overconfident on some predictions and underconfident on others, both deviations increase the MAE between cross-entropy and entropy. The fact that OKO achieves the lowest MAE in most settings suggests that its calibration improvement is not merely a bin-level artifact but reflects genuine per-prediction alignment between confidence and correctness.

---

Few-Shot Results: OKO Sets New State-of-the-Art on Limited-Data MNIST

The paper explicitly compares against Liu et al. (2022), which reported state-of-the-art few-shot MNIST results using transformation-invariant SVMs. The numbers, extracted from the main text (Section 5, "Generalization" paragraph):

• 10-shot MNIST (100 training points, 10 per class): OKO achieves 87.62% average accuracy, best seed 90.14%. Liu et al. (2022) best: 79.03%. Improvement: +8.59 percentage points.
• 20-shot MNIST (200 training points): OKO improvement of +2.85% over Liu et al. (2022)'s best.
• 50-shot MNIST (500 training points): OKO improvement of +1.81% over Liu et al. (2022)'s best.

These are substantial margins in a benchmark where state-of-the-art has been progressively improving. The fact that OKO — a general training framework with no MNIST-specific design — outperforms a method specifically designed for few-shot learning with transformation invariance is notable. However, the paper does not report few-shot results on the other datasets in a way that enables direct comparison to dedicated few-shot learning methods; the low-data curves in Figure 2 are not presented as formal few-shot benchmarks with per-class sample counts.

---

Ablation Studies and Robustness Checks

Number of odd classes $k$: The paper ablates $k \in \{0, 1, 2, 3\}$ in Appendix F.4 (Figures 11–13, Tables 5–6). The headline finding is that any $k \geq 1$ works, but $k = 0$ (pairs only, no odd class) fails. Figure 11 shows that $k=0$ ("Pair") achieves substantially lower test accuracy than $k \geq 1$ across all datasets and class distributions, with the gap being largest for heavy-tailed settings (e.g., MNIST heavy-tailed at 500 points: $k=1$ reaches ~88% while $k=0$ reaches ~80%). Figure 12 shows the same pattern for ECE: $k=0$ produces ECE values 2–5× higher than $k \geq 1$ in low-data regimes. Figure 13 and Table 5 confirm that $k=0$ produces dramatically worse RC alignment (MAEs 2–10× larger than $k=1$). Among $k \in \{1, 2, 3\}$, performance differences are relatively small: on uniform distributions, $k=1$ achieves slightly better accuracy and ECE; on heavy-tailed distributions, larger $k$ sometimes marginally improves RC alignment (Table 5: CIFAR-100 heavy-tailed MAE is 1.164 for $k=1$, 1.040 for $k=2$, 0.769 for $k=3$). The paper selects $k=1$ for all main experiments as the computationally cheapest option ($k+2=3$ examples per set, minimal overhead). Key insight: the presence of at least one contrastive (odd-class) example is essential — the calibration benefit is not simply from seeing multiple examples, but from seeing examples that force the model to distinguish between classes within each set.

Hard vs. soft loss: The paper ablates the loss function choice in Appendix F.5 (Figures 14–16, Table 6). The headline finding is that hard loss always outperforms soft loss in accuracy, and generally produces better calibration, with soft loss causing excessive underconfidence. Figure 14 shows that across all settings, hard loss achieves higher test accuracy than the corresponding soft loss variant, with gaps being larger for uniform distributions (MNIST uniform at 5000 points: hard ~97% vs. soft ~93%) than for heavy-tailed (MNIST heavy-tailed at 5000 points: hard ~96% vs. soft ~94%). Figure 15 shows that soft loss produces substantially worse ECE on MNIST and FashionMNIST (ECE of 0.30–0.60 vs. hard loss at 0.02–0.10), but the gap narrows on CIFAR-10 and disappears on CIFAR-100 (where soft loss sometimes yields slightly lower ECE). Figure 16 provides the mechanistic explanation: soft loss predictions show entropy concentrated near $\log C$ (maximum uncertainty) for both correct and incorrect predictions — the model is underconfident across the board because the soft target explicitly tells it to spread probability mass across all classes in the set. Table 6 quantifies this: the MAE between cross-entropy and entropy is 5–10× larger for soft loss than hard loss on MNIST/FashionMNIST. Key insight: explicitly encoding the set composition into the target distribution (soft loss) is counterproductive because it solves a proportion-estimation problem rather than a classification problem, producing a model that is calibrated but at the cost of being uniformly uncertain — it loses the ability to distinguish between cases where confidence is warranted and cases where it is not. The hard loss, by forcing the model to commit to a single answer, preserves the discriminative capacity while the set-level structure provides implicit regularization.

Auxiliary odd-class prediction head: The paper mentions in Section 5 that "preliminary experiments have shown that generalization performance can be boosted by predicting the odd class using an additional classification head." This auxiliary head is included in all reported experiments but is not ablated in the paper. Its contribution relative to the main OKO objective is unknown from the reported results — we cannot determine how much of OKO's accuracy advantage comes from the multi-task signal vs. the set-level objective itself. This is a notable gap in the ablation analysis.

Sensitivity to training data quantity: The paper's primary experimental design — varying the number of training points across a wide range — is itself an ablation of data quantity. The key finding (visible across Figures 2, 10, 11, 12): OKO's advantages over baselines are largest at small training set sizes and diminish as data increases. On MNIST uniform (Figure 2, top-left), the gap between OKO and vanilla shrinks from ~5 percentage points at 100 training points to ~1 percentage point at 5000 points. On FashionMNIST uniform, the gap shrinks from ~4 points at 200 to ~0.5 points at 10,000. This is consistent with the mechanism described in Section 4: in data-rich regimes, standard training has enough examples to learn appropriate confidence without needing the set-level implicit regularization. OKO's calibration benefit is specifically valuable when data is too scarce for standard training to learn calibrated probabilities on its own.

Architecture robustness: The paper uses different architectures for different datasets (CNN for MNIST/FashionMNIST, ResNet18 for CIFAR-10, ResNet34 for CIFAR-100) but does not ablate architecture within a single dataset. The consistent pattern of OKO outperforming baselines across all four architecture-dataset pairs provides suggestive evidence that OKO's benefits are architecture-independent, but this is not formally tested (e.g., running ResNet18 on MNIST to see if the magnitude of improvement changes).

---

Critical Assessment

Claim 1: "OKO simultaneously obtains better accuracy and calibration than standard training and existing methods, especially with class imbalance."

What was tested: The paper tests OKO against seven baselines on four datasets, two class distributions (uniform, heavy-tailed), and multiple training set sizes. Accuracy, ECE, reliability diagrams, and RC are all reported. OKO indeed achieves the best average accuracy in all settings (Table 2) and the lowest ECE in most settings (Figure 10), with the calibration advantage being most pronounced in low-data and heavy-tailed regimes.

What was not tested: (1) The paper does not evaluate OKO against modern calibration-training hybrids beyond label smoothing and focal loss — methods like Mixup (Thulasidasan et al., 2019, which the paper cites in related work but excludes from experiments), deep ensembles (Lakshminarayanan et al., 2017), or the soft calibration objectives of Karandikar et al. (2021) are not compared. (2) The auxiliary odd-class prediction head is not ablated, so we cannot separate its contribution from the core OKO objective. (3) All results are on small-to-medium image classification benchmarks; generalization to large-scale settings (ImageNet-scale), NLP tasks, or regression problems is untested. (4) Class-wise calibration metrics (Kull et al., 2019) are not reported — we only see aggregate ECE and RC, not whether OKO specifically improves minority-class calibration, which is the paper's motivation.

Assessment: The claim is well-supported for the tested benchmarks and baselines, with the caveat that the comparison set omits several relevant calibration methods and that the auxiliary head's contribution is unknown. The claim's qualifier — "especially with class imbalance" — is supported by the heavy-tailed results showing larger accuracy and ECE gaps than the uniform results, but the lack of class-wise calibration breakdown means we cannot verify that minority-class calibration specifically improves.

Claim 2: "OKO does not introduce additional hyperparameters for post-training tuning or require careful warping of the label distribution."

What was tested: OKO uses a single fixed $k=1$ for all experiments after ablating $k \in \{0,1,2,3\}$, uses the hard loss (selected after comparing against soft loss), and includes the auxiliary prediction head (selected after preliminary experiments). The baselines require tuning $\alpha$ for label smoothing or $\tau$ for temperature scaling.

What was not tested: (1) The choice of $k=1$ was made after observing that larger $k$ values perform similarly — but this observation comes from the same datasets used for evaluation, meaning $k=1$ is implicitly tuned on the test benchmarks. If the optimal $k$ were dataset-dependent and required tuning, OKO would introduce a hyperparameter. (2) The auxiliary head is a binary design choice (present/absent) whose impact is not quantified; if it provides a substantial accuracy boost, then OKO's performance depends on a design decision that might not transfer to other settings. (3) The number of OKO sets per epoch is set to $n_{\text{train}}$, which is a hyperparameter — the paper does not explore sensitivity to this choice (e.g., generating $2n_{\text{train}}$ or $n_{\text{train}}/2$ sets per epoch).

Assessment: The claim is partially supported. OKO indeed avoids the continuous hyperparameter tuning that label smoothing and temperature scaling require, but it substitutes this with discrete design choices ($k$, hard vs. soft loss, auxiliary head, sets per epoch) that are fixed after empirical evaluation on the test benchmarks. The paper's assertion that OKO is "hyperparameter-free" is an overstatement — it is more accurately described as having hyperparameters that are less sensitive and can be set to reasonable defaults ($k=1$, hard loss) without extensive tuning.

Claim 3: "OKO yields smoothed logits that result in accurate calibration, although models are trained using hard labels."

What was tested: The loss landscape analysis (Propositions 2–3) proves theoretical properties of the OKO risk that should produce finite logits. The entropy distribution plots (Figures 4, 8, 9) show that OKO models produce less extreme probabilities than vanilla training. The RC analysis (Figure 5, Table 3) confirms lower excess confidence.

What was not tested: (1) The paper does not directly report logit magnitudes or softmax temperature (a common diagnostic for overconfidence — if logits have large magnitude, softmax outputs are peaked). A direct comparison of average logit norms between OKO and vanilla training would provide stronger evidence for the "smoothed logits" claim. (2) The theoretical analysis assumes models have memorized the training data ($f_\theta(x_i) = F_{y_i}$), which is a useful simplification but may not capture the regularization effect during early training stages. The claim that OKO achieves calibration "naturally" without explicit smoothing penalties would be strengthened by showing that the logit smoothing emerges gradually during training, rather than being a property of the final converged state.

Assessment: The claim is theoretically grounded and empirically supported by indirect evidence (entropy distributions, ECE, RC), but direct evidence — logit magnitude comparisons, training dynamics of logit norms — is not provided, making the "smoothed logits" mechanism a reasonable inference rather than a proven fact.

Claim 4: "In few-shot settings, OKO achieves compellingly low calibration and classification errors" (Section 1 contributions).

What was tested: Few-shot MNIST accuracy is reported and compared against Liu et al. (2022). ECE in low-data regimes is visible in Figure 10 and shows OKO achieving lower ECE than baselines at small training set sizes.

What was not tested: (1) Formal few-shot evaluation (N-way K-shot episodes with held-out classes) is not performed — the low-data experiments use a fixed set of classes and vary the total number of training points, which is a different protocol from standard few-shot learning where the model must adapt to novel classes at test time. (2) Few-shot results are only reported for MNIST and only compared to one prior method; results on FashionMNIST, CIFAR-10, and CIFAR-100 at low data counts are shown in Figure 2 but not positioned as few-shot benchmarks with comparisons to dedicated few-shot methods. (3) Calibration in few-shot settings is not specifically analyzed — we see low-data ECE curves but no analysis of whether OKO's calibration advantage at 100 MNIST training points generalizes to other datasets at equivalent data scarcity.

Assessment: The few-shot claim is specifically supported for MNIST (where OKO achieves excellent results and beats a dedicated method) but is substantially weaker for the other datasets, where no few-shot benchmarking is performed. The "few-shot" framing is somewhat misleading — the experiments are better described as "low-data regime" experiments, since they train on a fixed label set with limited examples rather than evaluating generalization to unseen classes (the defining characteristic of few-shot learning).

Overall assessment of experimental design:

Strengths:

• Multi-faceted calibration evaluation (ECE, reliability diagrams, entropy distributions, RC) provides converging evidence rather than relying on a single metric.
• The heavy-tailed distribution setup (90/10 split with only 3 overrepresented classes) is severe and realistic, making the positive results more convincing.
• The data-quantity sweep (varying training points from extremely low to full dataset) directly tests the paper's motivation that OKO helps most when data is scarce.
• Five-seed error bands and cross-seed averaging address variance concerns.

Weaknesses:

• Missing baselines: Mixup (Thulasidasan et al., 2019), which is cited in related work as "most related to our approach" for calibration, is not evaluated. Deep ensembles, another standard calibration method, is absent. Temperature scaling without batch-balancing (vanilla + TS) is not reported, which would help isolate whether batch-balancing or temperature scaling drives the BB + TS results.
• Auxiliary head not ablated: The reported OKO performance includes an auxiliary odd-class prediction head that provides additional training signal. Without ablating it, we cannot attribute OKO's gains to the set-level objective vs. the multi-task learning effect. This is a significant confound.
• No class-wise breakdown: The paper motivates OKO by arguing that existing methods fail to calibrate minority classes specifically, but all reported calibration metrics (ECE, RC, reliability diagrams) are aggregate. Class-wise ECE or per-class RC analysis would directly test the motivating claim.
• Single-scale evaluation: All datasets are small-scale image classification. Transfer to larger benchmarks (ImageNet), different modalities (text, tabular), or different task types (regression, segmentation) is unknown.
• No computational cost analysis: The paper claims OKO "introduces little computational overhead," but at $k=1$, each training step processes 3× more examples. Total training time comparisons are not reported, making it impossible to evaluate whether OKO's accuracy/calibration gains are worth the additional compute.
• $k=1$ implicitly tuned: The choice of $k$ is ablated and found to matter little, but this ablation is itself a form of hyperparameter tuning on the test benchmarks. Whether $k=1$ generalizes as a default to new datasets is unknown.
• No confidence intervals on the RC analysis: Table 3 reports point estimates of MAE without any measure of variance, making it impossible to assess whether differences between methods (e.g., OKO at 0.334 vs. BB + LS at 0.114 on FashionMNIST heavy-tailed) are statistically significant or within noise.

Missing experiments that would strengthen the paper:

1. Ablation of the auxiliary odd-class head to isolate OKO's contribution.
2. Class-wise calibration metrics (ECE per class, especially minority vs. majority).
3. Direct logit magnitude or temperature analysis to validate the "smoothed logits" mechanism.
4. Comparison against Mixup and deep ensembles.
5. Wall-clock training time vs. accuracy/calibration tradeoff curves.
6. Evaluation on a larger-scale benchmark (e.g., ImageNet-LT for long-tailed recognition).
7. Formal few-shot learning protocol with held-out classes to validate the few-shot claim.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Dominates the Inference Budget and Is Excluded from All Efficiency Calculations

The assumption or constraint. The paper's headline results — including the 4× efficiency gains over best-of-N, the FLOPs-matched comparisons against the ~14× larger model, and the compute-optimal strategy selection — all assume that prompt difficulty is known before the inference budget is spent. The method used to estimate difficulty is computationally extreme: for each question, generate 2048 complete solutions from the base model and compute the pass@1 rate (oracle) or average PRM final-answer score (predicted). The authors acknowledge this explicitly in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

This is a candid admission, but it means the cost of the very mechanism that enables compute-optimal allocation is entirely excluded from the reported efficiency numbers.

The consequence. In a realistic deployment, the total cost is difficulty estimation + strategy execution, and the former can dominate the latter. Generating 2048 samples per prompt is equivalent to the largest test-time compute budgets studied in the paper (256–512 generations for the strategy itself). For a prompt that turns out to be easy — where the compute-optimal strategy might allocate only 4–8 generations — the difficulty estimation step would consume 250–500× more compute than the actual solution process. This means the reported 4× efficiency gains over best-of-N are computed after difficulty is already known, without amortizing the cost of learning it. Unless difficulty can be estimated far more cheaply (e.g., with a lightweight classifier or from a small number of initial samples), the practical efficiency of the compute-optimal framework may be substantially lower than the paper's figures suggest — potentially worse than a uniform best-of-N baseline when estimation overhead is included.

What evidence exists in the paper. The only evidence that difficulty estimation might be made cheaper is the finding that predicted difficulty bins (using PRM scores) track oracle bins closely — the curves "largely overlap" in Figures 4 and 8. However, the predicted method still requires 2048 samples per question plus PRM scoring, which is only cheaper than the oracle in that it removes the need for ground-truth labels — it does not reduce the computational cost. The paper does not report any experiment with a cheaper difficulty estimator, nor does it include difficulty estimation overhead in any budget calculation. Section 8 explicitly flags this as a key direction for future work:

"estimating difficulty in this way still incurs additional computation cost during inference... we leave the exploration of more sophisticated techniques for predicting question difficulty to future work"

Mitigation status. Not addressed. The paper frames the difficulty estimation cost as an exploration-exploitation tradeoff and suggests future work on "pretraining or finetuning models to directly predict difficulty of a question" (Section 8) and "adaptive difficulty estimation" where initial samples inform subsequent allocation. But no such model or adaptive scheme is developed or evaluated. Until this gap is closed, the reported efficiency gains represent an upper bound on achievable deployment efficiency, not a realized gain.

---

Hard Problems Are Fundamentally Unsolved — Test-Time Compute Creates No New Capability

The assumption or constraint. The paper's framework assumes that the base model's proposal distribution contains correct solutions at some non-trivial rate — i.e., that pass@1 > 0 for the question. This is explicit in the difficulty definition (Section 3.2), where difficulty is measured by the base model's pass@1 rate, and in the FLOPs-matched analysis takeaway:

"test-time compute can amplify existing capability but does not create it from nothing"

For the hardest questions (difficulty bin 5, pass@1 near zero), no amount of search, revision, or compute-optimal allocation produces meaningful improvement.

The consequence. Across all methods — search, revisions, and their compute-optimal combinations — bin 5 accuracy hovers at approximately 1–3% regardless of compute budget (Figure 3, right; Figure 7, right). In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5% for both revisions and PRM search, and test-time compute with the smaller model is substantially worse than the ~14× larger model at all values of $R$ (e.g., hard questions show a −52.9% relative disadvantage from using test-time compute instead of the larger model in the PRM search comparison at $R \gg 1$). This means that for genuinely novel, out-of-distribution, or capability-stretching problems, pretraining is the only viable path to improvement — test-time compute offers essentially zero benefit. The compute-optimal framework provides no mechanism to address this fundamental ceiling.

What evidence exists in the paper. Every difficulty-binned analysis in the paper shows this pattern. Figure 3 (right): bin 5 accuracy is flat at ~1–3% across all budgets and methods. Figure 7 (right): bin 5 accuracy is flat at ~2–3% across all sequential-to-parallel ratios. Figure 9: bin 5 scaling curves (blue, bottommost) are essentially horizontal lines at 0–5% accuracy for both revisions and PRM search, well below the larger model's greedy performance (shown as stars). Table 2 in Section 7: hard questions show negative or near-zero relative improvement from test-time compute vs. the larger model across almost all $R$ values. The evidence is consistent and unambiguous.

Mitigation status. The paper is transparent about this limitation, explicitly stating it in Section 7 and the conclusion. However, it offers no mitigation: the boundary is fundamental to the approach. Test-time compute can only select or refine solutions the base model can already produce at some rate; if the base model's pass@1 is zero, no search or revision strategy can recover correct answers. This is not a flaw in the method but a hard capability bound that practitioners must account for when deciding between investing in larger pretrained models vs. better test-time strategies.

---

Single Benchmark and Single Model Family — No Evidence of Generalization Across Domains or Architectures

The assumption or constraint. All experiments are conducted on the MATH benchmark (500 test questions, high-school competition-level math) using PaLM 2-S* (Codey) as the base model. The authors state they "believe this model is representative of the capabilities of many contemporary LLMs" (Section 4), but this claim is not empirically tested. The choice of MATH is deliberate — the authors argue test-time compute "is expected to help most when the model already possesses the necessary knowledge and the challenge is drawing complex inferences" (Section 4) — but this means the findings are specifically calibrated to symbolic math reasoning and may not transfer to other task families.

The consequence. Several aspects of the findings could be model-specific or domain-specific. The PRM's quality and over-optimization behavior — which governs when beam search helps vs. hurts — depend on PaLM 2-S*'s output distribution, calibration, and error patterns. A model with different properties (e.g., better calibrated base outputs, different step-level reasoning structure) might exhibit different difficulty-dependent scaling curves that would change the compute-optimal policy. The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning and self-correction capabilities, which vary substantially across model families. The MATH benchmark consists of problems with well-defined symbolic answers graded by exact string matching — it is unclear whether the difficulty-dependent patterns (beam search hurting easy problems, revisions helping easy problems) generalize to code generation (where unit tests provide correctness signals), logical reasoning, scientific QA, or tasks requiring factual recall rather than multi-step deduction. Domains without clean correctness signals (open-ended generation, summarization, dialogue) would require fundamentally different verifier training and difficulty estimation approaches that the paper does not address.

What evidence exists in the paper. None — this limitation is purely about scope. The paper does not include experiments on any benchmark other than MATH, nor does it test any model other than PaLM 2-S*. The ablation studies (PRM aggregation, revision verifier, oracle vs. predicted bins) are all conducted within the MATH + PaLM 2-S* setting. There is no evaluation on code generation (HumanEval, MBPP), commonsense reasoning, scientific QA, or any non-math domain. There is no evaluation with a different base model family (e.g., LLaMA, GPT, Gemini) to test whether the key findings — verifier over-optimization thresholds, revision effectiveness per difficulty tier, compute-optimal strategy patterns — replicate.

Mitigation status. The authors acknowledge the single-benchmark limitation implicitly by not claiming broader generalization, but they do not explicitly flag it as a limitation in Section 8 (Limitations and Future Directions in the paper) or discuss what would be needed to validate transfer to other domains. The paper's positioning as a systematic scaling analysis rather than a new method means that single-benchmark evaluation is more acceptable — the contribution is the analytical framework, not a claim of universal performance — but practitioners considering deploying these techniques on non-math tasks have no empirical guidance from the paper on whether to expect similar benefits.

---

Verifier Over-Optimization Is the Hard Ceiling on Scaling, and the Paper Provides No Robustness Solution

The assumption or constraint. The compute-optimal framework routes easy problems away from aggressive optimization (beam search) because the PRM's scores become unreliable under strong optimization pressure. On easy problems (bins 1–2), beam search degrades performance with increasing budget (Figure 3, right) — the model finds solutions that score highly under the PRM but are incorrect, a phenomenon the paper calls "over-optimization of the PRM." This is not a minor edge case: it is the primary bottleneck preventing unbounded improvements from increased test-time compute. The compute-optimal policy mitigates this by routing easy problems to best-of-N and only applying beam search on medium-difficulty problems, but it does not solve the underlying problem — verifier reliability remains the ceiling.

The consequence. Even with compute-optimal allocation, the benefits of scaling test-time compute are fundamentally capped by verifier quality. On medium-difficulty problems where beam search is deployed (bin 3–4), the beam search curves in Figure 3 flatten and sometimes decline well before the maximum budget is exhausted — evidence that over-optimization limits even the "optimal" strategy. Lookahead search — the most powerful optimizer, which simulates additional steps forward to improve step-level scoring — paradoxically performs worst overall at the same generation budget (Figure 3, left) because its extra computational cost reduces the effective number of beams explored, and the improved scoring still falls prey to over-optimization. Qualitative examples in Appendix M (Figure 29) show search producing degenerate outputs (repetitive low-information steps at the end of solutions, overly short 1–2 step solutions) that score highly under the PRM but are nonsense. This means that even an optimally allocated test-time compute budget cannot overcome a weak verifier — the quality of the verifier is the ultimate determinant of how much test-time compute helps. This shifts the research bottleneck from "how do we allocate compute?" to "how do we build robust verifiers?", but the paper explores no approaches for improving verifier robustness beyond the initial Monte Carlo rollout training procedure.

What evidence exists in the paper. Figure 3 (right): beam search accuracy on bin 1 decreases from ~78% to ~77% as budget increases from 4 to 256, while best-of-N increases from ~68% to ~88%. Figure 3 (left): lookahead search underperforms all methods at equivalent budgets, despite being the most sophisticated optimizer. Appendix M: qualitative examples of degenerate beam search outputs (repetitive steps, overly short solutions). The paper explicitly names this phenomenon in Section 5.3:

"The degradation at high budgets is attributed to over-optimization of the PRM — search finds solutions that score highly under the PRM but are actually incorrect."

The paper does not report a systematic analysis of why the PRM over-optimizes on certain problems, what properties of the training data or architecture contribute to it, or how it might be mitigated beyond the compute-optimal routing strategy (which treats the symptom, not the cause).

Mitigation status. Partially addressed through compute-optimal routing (the meta-strategy of avoiding aggressive optimization where the verifier is unreliable), which helps but does not solve the problem. The paper does not explore adversarial PRM training (training on search-generated solutions rather than i.i.d. samples), ensemble verification, or constrained search with KL penalties relative to the base model's output distribution. Section 8 flags verifier robustness as a key future direction: the over-optimization finding "redirects research attention: rather than developing ever-more-sophisticated search algorithms... the priority should be building more robust verifiers." This is an accurate diagnosis, but it means the current method's scaling ceiling is determined by verifier quality, and practitioners adopting OKO inherit whatever limitations their verifier has.

---

The Test Set Is Small (500 Questions, Split Into Quintiles of ~100 Each) — Strategy Selection Has High Variance and the Policy May Not Be Robust

The assumption or constraint. The compute-optimal policy is selected via two-fold cross-validation on the 500-question MATH test set, split into five difficulty quintiles of approximately 100 questions each. Each fold contains roughly 50 questions per difficulty bin. The best-performing strategy (search algorithm, revision depth, sequential-to-parallel ratio) is selected on one fold and evaluated on the other, with results averaged. This means the strategy selection — which determines the entire compute-optimal behavior — is based on performance measured on ~50 questions per bin.

The consequence. Strategy selection on such small sample sizes is high-variance. A different random split of the 500 questions could select different optimal strategies per bin, especially for settings where multiple strategies perform similarly. The paper does not report confidence intervals on the compute-optimal scaling curves (Figures 4 and 8 show point estimates without error bands) or test the sensitivity of the selected policy to the cross-validation split. This makes it impossible to assess whether the reported 4× efficiency gains are statistically reliable or could vary substantially with a different test set or split. For medium-difficulty problems (bins 3–4), where the gap between beam search and best-of-N is relatively small (a few percentage points in Figure 3, right), the selection of beam search as the "optimal" strategy could easily flip with a different sample of 50 questions. If the selected strategies are not robust to the evaluation sample, a practitioner deploying the compute-optimal policy on a new problem distribution could experience performance degradation rather than the reported gains.

What evidence exists in the paper. The paper reports that oracle and predicted difficulty bins produce similar compute-optimal curves (Figures 4 and 8), which provides some robustness evidence — the strategy selection is not highly sensitive to whether difficulty is measured via ground-truth pass@1 or PRM score. However, this does not address variance due to the small evaluation sample itself. The 95% confidence intervals shown in other figures (e.g., Figure 2 for accuracy, Figure 10 for ECE) are computed over five random seeds of model training, not over different test set splits or different cross-validation folds. The compute-optimal scaling curves (Figures 4, 8) show no error bands at all. The paper does not report ablation of the number of folds (e.g., comparing 2-fold vs. 5-fold vs. leave-one-out cross-validation) or the sensitivity of selected strategies to the fold assignment. The paper does not report how often the same strategy is selected across folds — a direct measure of policy stability.

Mitigation status. Not addressed. The two-fold cross-validation protocol is described once in Section 3.2 with no further analysis of its statistical properties. The paper does not discuss the sample size limitation, does not report variance estimates for the compute-optimal policy, and does not explore alternative policy selection methods (e.g., fitting a parametric function that maps difficulty to strategy, which could be more sample-efficient than per-bin lookups). This is particularly concerning given that the test set size (500 questions) is fixed by the MATH benchmark and cannot be increased without changing the evaluation protocol.

---

The Revision Model Exhibits a 38% Correct-to-Incorrect Reversion Rate, and Revision Training Is Fragile to Data Generation Choices

The assumption or constraint. The revision model is fine-tuned exclusively on sequences where all in-context answers are incorrect, followed by a correct answer. During supervised fine-tuning, the model never sees correct answers in the context, so it never learns what to do when the current answer is already correct. At test time, when the model produces a correct answer during a revision chain, the next revision step may incorrectly "revise" it into a wrong answer because the training data provided no signal for this case. The paper reports (Section 6.1):

"approximately 38% of correct answers get converted back to incorrect ones using a naive approach"

Furthermore, the reversal of gains in the ReST$^{\text{EM}}$ experiment (Appendix K, Figure 16) — where additional sequential revisions substantially hurt performance with the RL-fine-tuned model — demonstrates that revision training is fragile to the data generation methodology. The ReST$^{\text{EM}}$-trained model, which used on-policy data collection, achieved approximately 33.5% accuracy with fully sequential revisions at 256 generations compared to roughly 38.5% at the optimal ratio, whereas the SFT-trained model improved monotonically with sequential revisions.

The consequence. The 38% reversion rate means that longer revision chains are not monotonically beneficial — at each step, there is a substantial probability of degrading a correct answer. The paper's mitigation — using majority voting or verifier-based selection across the entire chain rather than always taking the last revision — is an imperfect patch. It recovers some of the lost correct answers but does not prevent the model from wasting computation generating invalid revisions from correct starting points. In latency-constrained settings where generating a long chain of revisions is expensive, the reversion problem means that later revisions are not only unhelpful but actively harmful, requiring the selection mechanism to correctly identify and reject them. The ReST$^{\text{EM}}$ failure demonstrates that the positive results depend on specific training data construction choices (offline data with edit-distance-based incorrect-correct pairing) that may not transfer to other settings or to iterative self-improvement pipelines where on-policy data generation is natural.

What evidence exists in the paper. The 38% reversion rate is reported explicitly in Section 6.1 (the paper's description of the correct-to-incorrect reversion problem). The ReST$^{\text{EM}}$ degradation is shown in Appendix K, Figure 16, where the ReST$^{\text{EM}}$-trained model's fully sequential performance drops sharply. The paper also notes (Section 6.1) that standard validation loss is not a reliable signal for early stopping during revision model training because "after fine-tuning, the validation trajectories become off-policy (they were generated by the base model, not the fine-tuned revision model)" — requiring heuristic early-stopping based on the point where validation loss begins increasing. This off-policy issue is a structural problem with revision training that the paper identifies but does not resolve.

Mitigation status. The paper provides a partial mitigation (within-chain selection via majority voting or verifier) that reduces the impact of the reversion problem but does not eliminate it — correct answers that get incorrectly revised are recovered only if the selection mechanism can identify them, and the selection mechanism itself has error. The paper does not explore more principled solutions, such as training the revision model on trajectories that include correct answers in context with a "stop revising" signal, or using the PRM to decide when to terminate the revision chain early. The ReST$^{\text{EM}}$ failure is presented as a negative result but is not analyzed in depth — the authors hypothesize that "on-policy data collection in ReST$^{\text{EM}}$ exacerbates spurious correlations in revision data" (Appendix K) but do not investigate what specific correlations cause the degradation or how to avoid them. For practitioners seeking to deploy OKO-style revisions, the message is that the specific offline data construction procedure matters critically, and deviating from it (e.g., to use on-policy data for iterative improvement) can cause substantial performance regression.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a conceptual reframing rather than a paradigm shift: it redefines calibration from a post-hoc correction or loss-function modification into a training data structure problem. The central move — training on constructed sets where the model must identify the majority class from aggregated logits — demonstrates that calibration-like regularization can emerge from the data presentation itself, without explicit calibration penalties, hyperparameter tuning, or post-processing. This reframing opens a new design axis for the field: rather than asking "what loss function or post-processing produces well-calibrated probabilities?", researchers can ask "what set construction induces the right calibration properties for a given task?"

The magnitude of this shift is incremental but practically meaningful. OKO does not displace the need for calibration-aware training — it provides a new mechanism that complements existing approaches (label smoothing, temperature scaling, focal loss) rather than rendering them obsolete. However, it resolves a tension that prior work accepted as inevitable: that improving calibration requires either sacrificing accuracy (weight decay, which the paper notes "has been shown to improve calibration" but "comes at a cost of generalization"; Guo et al., 2017, cited in Appendix C), imposing rigid structures on logits (label smoothing forces convergence to a fixed target distribution; Proposition 3 and Appendix C discussion), or adding post-training complexity (temperature scaling requires a held-out calibration dataset). OKO achieves simultaneous accuracy and calibration improvements across all tested benchmarks and class distributions (Table 2, Figure 1B), demonstrating that this tradeoff is not fundamental but rather an artifact of single-example training.

The paper also reconciles a subtle contradiction in the calibration literature. Label smoothing is known to improve calibration (Müller et al., 2019) but the mechanism — replacing hard targets with soft targets — seemed necessary: if hard labels cause logit divergence and overconfidence, then softer targets should be required to prevent it. OKO demonstrates that hard labels can produce well-calibrated models when the training data is structured as sets rather than individual examples. The "smoothed logits" emerge from the optimization dynamics of the set-level objective (Propositions 2–3, Theorem 1), not from target modification. This suggests that the root cause of overconfidence is not hard labels per se, but the absence of between-example coupling during training — when the model must reason about sets of examples, the hard label signal is balanced by the need to account for all examples in the set, producing implicit regularization. This finding narrows the search space for calibration methods: future work need not choose between hard and soft targets, but can instead focus on how to structure training episodes to provide the right between-example couplings.

The finding that odd-class examples are essential ($k=0$ fails dramatically; Appendix F.4, Figures 11–13) while the exact number of odd classes matters little ($k=1,2,3$ perform similarly) points toward a specific mechanism: the calibration benefit comes from contrastive pressure — the model must distinguish the pair class from explicitly present alternatives within each training episode. This connects OKO conceptually to contrastive representation learning (where negative examples are essential) and suggests that the design space for set-based training is richer than simply "train on sets." The composition of the set (same-class pairs vs. odd-class distractors, the ratio of pair to odd examples, the relationship between odd class selection and class frequencies) determines what regularization properties emerge. The paper's specific construction — uniform pair class selection (implicit class balancing), two within-class examples, and $k$ distinct between-class examples — is one point in this design space that happens to work well, but the ablation results indicate that $k$ is not the critical parameter; the presence of contrastive examples is.

The paper redirects research attention in two specific ways:

• Toward training data structure as a first-class design element for calibration. Prior work focused on loss functions (focal loss, label smoothing, proper scoring rules), post-processing (temperature scaling, isotonic regression, Dirichlet calibration), or architecture (ensembles, Monte Carlo dropout). OKO demonstrates that simply changing how examples are grouped during training — without modifying the loss, architecture, or post-processing — can produce calibration benefits comparable to or better than these approaches. This opens a research direction that had not been systematically explored: what other set constructions (triplets, hierarchical sets, adaptive set composition based on model uncertainty) yield desirable calibration or robustness properties?

• Away from hyperparameter-intensive calibration methods for low-data settings. Label smoothing requires tuning $\alpha$; temperature scaling requires tuning $\tau$ on a held-out set; focal loss requires tuning the focusing parameter $\gamma$. In low-data regimes, held-out calibration data may be insufficient for reliable tuning, making these methods brittle. OKO's approach — fix $k=1$, use the hard loss, and optionally include an auxiliary head — requires no continuous hyperparameter search, making it practically deployable in exactly the data-scarce settings where calibration matters most (medical imaging, scientific data, rare event prediction). The paper's strong few-shot MNIST results (8.59% improvement over prior best at 10-shot; Section 5) demonstrate this concretely. This shifts the conversation from "which calibration method has the bestheld-out-tuned performance?" to "which method works robustly without tuning when data is too scarce for a calibration set?"

The paper does not, however, demonstrate that OKO's specific set construction is optimal or even near-optimal. The ablation of $k$ (Appendix F.4) shows insensitivity to $k \in \{1,2,3\}$ but does not explore other dimensions of the design space: What happens when the number of pair-class examples is larger than 2? What if odd classes are sampled with probability proportional to their frequency rather than uniformly? What if the set composition is adaptive based on model confidence during training? These questions remain open, and the primary landscape change is that the field now has a concrete, working example that such questions are worth asking.

Follow-Up Research This Work Enables

1. Adaptive set composition during training based on model uncertainty. The paper constructs OKO sets with a fixed structure (always 2 pair examples, $k$ odd examples, uniformly selected pair class). Theorem 1 shows that this uniform structure produces epistemic uncertainty calibration — the model is uncertain about sparse regions. But what if the set composition were adaptive? For example: identify which classes the model is currently overconfident about (using RC or per-class ECE on a validation set), and increase the frequency with which those classes appear as odd examples in sets where the pair class is well-represented. This would provide targeted contrastive pressure where overconfidence is worst. A concrete experiment: on CIFAR-100 heavy-tailed, track per-class RC during training and dynamically adjust the probability of selecting each class as an odd class inversely to its current RC — does this produce better minority-class calibration than the static uniform selection? The paper's RC measure (Definition 1) provides the per-datapoint tool needed to implement this adaptation, and the finding that odd-class presence is crucial (Appendix F.4) but the number $k$ is not suggests that which classes serve as odds matters more than how many.

2. OKO for regression and structured prediction. The paper frames OKO entirely within classification, but the core mechanism — sum the outputs for a set, compare against a target derived from the set composition — could extend to regression. For example: construct sets where two inputs have the same (or similar) target value and $k$ inputs have different target values; the model predicts the sum of outputs across the set, with the loss comparing against the sum of true targets. This would force the model to recognize which inputs share similar regression targets, implicitly regularizing predictions in sparse regions analogously to Theorem 1. A concrete first experiment: UCI regression benchmarks with heavy-tailed target distributions (e.g., Boston Housing with a few very high-value properties, or medical cost prediction with a few very expensive patients). Train an OKO-style regression model (sum outputs, MSE loss against summed targets) and compare calibration (using quantile calibration metrics rather than ECE) to standard MSE, quantile regression, and Gaussian process baselines. The key question: does set-based regression show the same "epistemic uncertainty manifested as predictive variance" pattern that Theorem 1 proves for classification?

3. Class-wise calibration analysis to validate the minority-class calibration claim. The paper motivates OKO by arguing that existing imbalance methods "do not improve calibration for minority instances" (Section 1) and that "standard calibration procedures tend to systematically underestimate the probabilities for minority class instances" (citing Wallace & Dahabreh, 2012). However, all reported calibration metrics (ECE, RC, reliability diagrams) are aggregate — they do not separate majority-class from minority-class calibration. A direct follow-up would compute class-wise ECE for OKO vs. batch-balancing + label smoothing on the heavy-tailed CIFAR-100 setting, binning predictions by class frequency (overrepresented 3 classes vs. underrepresented 97 classes). The paper's entropy distribution plots (Figure 4) hint at better tail-class calibration — OKO shows more entropy spread for correct tail predictions, indicating appropriate caution — but this is qualitative. Numbers are needed: does OKO achieve lower class-wise ECE specifically for the 97 minority classes, or does its aggregate ECE advantage come primarily from fixing overconfidence on majority classes? If the latter, the motivating claim is unsupported and a different set construction might be needed for minority-class calibration.

4. Combining OKO with Mixup: do set-level and interpolation-based regularization compose? The paper cites Mixup (Thulasidasan et al., 2019) in related work as "most related to our approach" for calibration through data augmentation, noting that Mixup "blend[s] different inputs together" — but OKO and Mixup work through completely different mechanisms (set-level logit aggregation vs. convex combination of inputs and labels). They are potentially complementary: Mixup creates new training examples through interpolation; OKO restructures how existing examples are presented. A concrete experiment: train ResNet18 on CIFAR-10 and CIFAR-100 with (a) standard training, (b) Mixup alone, (c) OKO alone, (d) OKO + Mixup (apply Mixup to individual examples before grouping them into OKO sets). Hypotheses: OKO + Mixup might outperform either alone if the mechanisms are complementary (Mixup handles between-class interpolation uncertainty, OKO handles within-class frequency calibration), or they might interfere if Mixup's interpolated labels confuse the OKO set-level target (which assumes discrete, clean class assignments). The heavy-tailed setting is the most interesting stress test: Mixup is known to help with minority class generalization; OKO helps with calibration. Do they compose multiplicatively or does one dominate? The paper's code release makes this experiment straightforward to implement.

5. Training a lightweight difficulty predictor for OKO-style set sampling in active learning. The paper's Theorem 1 shows that OKO produces uncertainty calibrated to data frequency — the model is uncertain in regions with few samples. This property is exactly what active learning query strategies aim to measure: which unlabeled examples would reduce epistemic uncertainty most if labeled? A concrete follow-up: use an OKO-trained model's predictive entropy or RC score as an acquisition function for active learning. On CIFAR-10, start with 100 random labeled examples, train with OKO, then query the 10 unlabeled examples with highest RC or predictive entropy for labeling, retrain, and repeat. Compare against standard acquisition functions (entropy, margin, BALD) using a vanilla-trained model. Hypothesis: OKO's built-in epistemic uncertainty quantification (epistemic → aleatoric conversion) might provide better acquisition signals than a vanilla model's overconfident predictions, especially early in active learning when labeled data is extremely scarce. This would connect OKO directly to a downstream application where its properties provide a natural advantage, and would test whether the "uncertainty in sparse regions" property from Theorem 1 translates to practical active learning gains.

6. Stress-testing the set composition design space: what breaks OKO? The paper shows that $k=0$ (no odd class) breaks OKO and $k \geq 1$ works, but many other dimensions of set composition are unexplored. Specific stress-test experiments:

• Pair class frequency imbalance: Sample the pair class proportional to its training set frequency rather than uniformly. Hypothesis: this should destroy OKO's class-balancing effect and degrade minority-class accuracy to batch-balancing baseline levels. This would isolate whether the uniform pair class sampling or the set-level aggregation drives the gains.
• Weak odd examples: Replace the $k$ odd-class examples with examples from classes that are visually or semantically similar to the pair class (e.g., use a confusion matrix from a pretrained model to select "hard negative" odd classes). Hypothesis: harder odd examples may force better decision boundaries and further improve calibration, or may cause the model to become underconfident on the pair class.
• Multiple pair examples ($>2$): Use sets with 3, 4, or 5 examples from the pair class and $k$ odd examples. Hypothesis: increasing the within-class multiplicity should strengthen the pair-class signal and potentially improve calibration further, but at increased computational cost.
• Varying $k$ with dataset size: The paper found $k$ insensitive on the tested benchmarks, but on datasets with very many classes (e.g., ImageNet with 1000 classes), larger $k$ may be needed to provide sufficient contrastive pressure. Test OKO on ImageNet-100 (a 100-class subset) with $k=1,5,10,20$ to see if the insensitivity holds at scale.

Practical Applications and Downstream Use Cases

1. Low-data scientific and medical image classification. The paper's strongest results are in low-data regimes: on 10-shot MNIST (100 total training points), OKO achieves 87.62% accuracy with 90.14% for the best seed — an 8.59 percentage point improvement over the prior state-of-the-art few-shot method (Liu et al., 2022), while simultaneously achieving better calibration (lower ECE at 100 training points in Figure 10, top-left). These are exactly the conditions faced in many scientific and medical domains: limited labeled data (expensive expert annotation), class imbalance (common diseases vs. rare conditions), and high stakes where overconfidence on incorrect predictions is dangerous. A pathology lab with 50 labeled examples of a rare cancer subtype and 500 examples of common subtypes could train an OKO classifier that (a) achieves higher accuracy on the rare subtype than standard training or batch-balancing (analogous to the heavy-tailed MNIST results where OKO outperforms BB by ~7 percentage points at 500 training points), and (b) produces calibrated probabilities that indicate appropriate uncertainty about rare-subtype predictions (analogous to the entropy distribution shift in Figure 4 where OKO shows more entropy spread for tail-class correct predictions). The fact that OKO requires no held-out calibration set (unlike temperature scaling) is critical here — with only 50 rare examples, there is no data to spare for tuning $\tau$.

2. Deploying calibrated classifiers on edge devices without post-processing pipelines. The paper emphasizes that OKO models are applied to single examples at test time "exactly like any network trained via single-example learning" (Section 1) with "zero changes at inference time" (Appendix A). For on-device deployment (mobile phones, embedded systems, IoT), this matters: a standard calibration pipeline of training + temperature scaling requires storing and applying the temperature parameter at inference, which is trivial but adds a step to the deployment pipeline. More importantly, if the deployment data distribution shifts relative to the calibration set, the temperature parameter becomes stale and recalibration requires on-device held-out data that may not exist. An OKO-trained model embeds calibration into the weights themselves — the logits are naturally smoothed through the training process (Proposition 2) rather than scaled post-hoc. A practitioner training an OKO classifier for on-device image recognition (e.g., species identification from camera trap photos, where training data is limited and class distribution is heavy-tailed toward common species) can deploy the model directly with no calibration step, confident that the output probabilities reflect appropriate uncertainty for rare species and low-data regions. The entropy distribution plots (Figures 4, 8, 9) provide the visual evidence: OKO corrections on rare classes show appropriate uncertainty (entropy shifted away from 0), which vanilla training does not provide.

3. Batch inference pipelines where calibration quality affects downstream decision thresholds. In applications where model predictions feed into automated decision systems — loan approval, content moderation, triage — the predicted probabilities are used to set thresholds: "approve if P(repay) > 0.8," "flag for review if P(toxic) > 0.5." Miscalibrated probabilities directly cause misallocation: an overconfident model sets thresholds that are too permissive (approving loans that should be rejected because the model's 0.8 confidence doesn't correspond to 80% repayment rate). The paper's RC analysis (Figure 5, Table 3) shows that OKO models have lower mean absolute deviation between cross-entropy and entropy — meaning their probabilities are more trustworthy for thresholding. In a batch inference setting (e.g., nightly processing of loan applications), switching from a vanilla-trained classifier to an OKO-trained classifier (with no other changes to architecture or post-processing) would produce probability estimates where 0.8 confidence more reliably corresponds to ~80% accuracy, reducing the need for manual threshold calibration on held-out data. The quantitative backing: on FashionMNIST uniform, OKO achieves RC MAE of 0.080 vs. 0.455 for vanilla and 0.243 for vanilla + label smoothing (Table 3) — the OKO model's probabilities are substantially better aligned with its actual error rates, which directly translates to more reliable threshold-based decisions.

4. Self-training and pseudo-labeling with calibrated confidence estimates. Self-training pipelines iteratively label unlabeled data using the model's own predictions, selecting examples where the model is confident. Overconfident models produce overconfident pseudo-labels — they assign high confidence to incorrect predictions, which then get added to the training set as "ground truth," compounding errors. OKO's calibration properties (lower ECE, lower RC, appropriate uncertainty on sparse classes) suggest it would produce more reliable pseudo-labels. A concrete deployment: on a dataset where only 10% of examples are labeled (mimicking the heavy-tailed class distribution used in the paper), train an OKO classifier on the labeled set, use it to pseudo-label the unlabeled data with a confidence threshold, add high-confidence pseudo-labels to the training set, and retrain. Compare against the same pipeline using a vanilla-trained classifier. Hypothesis: OKO's lower overconfidence should mean fewer incorrect pseudo-labels above any given confidence threshold, leading to cleaner self-training iterations and better final accuracy. The paper's heavy-tailed results (OKO maintains accuracy on minority classes where vanilla training collapses; Figure 2, bottom row) provide suggestive evidence for this benefit.

When to Prefer This Method

The paper explicitly positions OKO against standard cross-entropy training, label smoothing, and class imbalance methods (focal loss, error re-weighting, batch-balancing). It does not explicitly position OKO against Mixup, deep ensembles, or Bayesian methods — so a general "Prefer OKO when..." matrix would impute claims the paper doesn't make. However, the paper does articulate clear conditions under which OKO's advantages manifest, based on its experimental results and theoretical analyses. These can be framed as a decision rule:

Prefer OKO when:

• Training data is limited (few-shot to low hundreds of examples per class): OKO's accuracy and calibration advantages are largest at small training set sizes (8.59% improvement over prior best 10-shot MNIST; Figures 2 and 10 show gaps narrowing as data increases). In data-rich regimes, vanilla training with label smoothing or temperature scaling performs comparably.
• Class distribution is heavy-tailed or imbalanced: OKO's strongest gains occur in heavy-tailed settings (Table 2: +4.25% over best baseline on MNIST heavy-tailed vs. +1.28% on uniform; ECE gaps are larger in heavy-tailed rows of Figure 10). The uniform pair class sampling in Algorithm 1 provides implicit class balancing without the probabilistic biases that direct undersampling introduces (Dal Pozzolo et al., 2015a, cited in Section 1).
• A held-out calibration set is unavailable or too small for reliable temperature tuning: OKO requires no post-hoc calibration and introduces no continuous hyperparameters needing grid search (the paper fixes $k=1$ after ablation). This makes it deployable in settings where temperature scaling or isotonic regression are infeasible due to data scarcity.
• Model architecture cannot be changed and post-processing complexity is undesirable: OKO modifies only the training data presentation and loss computation — the model architecture, inference code, and deployment pipeline are identical to standard classification. For practitioners with fixed infrastructure who want calibration without adding temperature scaling or ensemble logic to their serving stack, OKO is a drop-in training procedure change.

Prefer label smoothing or temperature scaling instead when:

• Training data is abundant and a held-out calibration set exists: At large training set sizes (e.g., full CIFAR-100, Figure 10, top-right), OKO's ECE is comparable to or slightly worse than batch-balancing + label smoothing. If a calibration set is available, tuning $\tau$ or $\alpha$ is cheap and effective; the marginal benefit of switching to OKO may not justify changing the training pipeline.
• Computational cost of 3× per-step training is prohibitive: With $k=1$, each OKO training step processes 3 examples (2 pair + 1 odd) instead of 1, increasing per-epoch computation by approximately 3× (the paper does not report wall-clock times, but this follows from the set size). If training throughput is the bottleneck, standard training with post-hoc calibration may be preferable.
• The class structure has meaningful semantic distance: OKO treats all pairs of classes as equally distant (Section 6, "one caveat of OKO is that classes are treated as semantically equally distant"). If the application benefits from label smoothing with class-conditional smoothing parameters (e.g., confusing similar classes more than distant ones), OKO's uniform treatment of odd classes may be suboptimal.