MNIST Loss Bounds per Example and an Adaptive Calibration Error (ACE) Implementation

An interview for a research scientist role included a machine-learning round built around a classifier for the MNIST handwritten-digit dataset, which has 10 classes. The round asked for loss bounds for a single example, an implementation of the adaptive calibration error (ACE), test cases for that implementation, and an explanation of why the ACE curve is spiky. Assume the model outputs a softmax probability vector $p = (p_1, \dots, p_{10})$ over the digits and is trained with cross-entropy loss; the report does not name the loss, so confirm it first. NumPy is allowed but not required. ### Clarifying Questions - Is the loss cross-entropy on the softmax output, or something else such as squared error? - When two classes tie for the highest probability, does the prediction count as correct or incorrect? - Should ACE use the probability of every class, as in the definition below, or only the probability of the predicted class? - How many bins should be used, and how should the predictions be split when their number is not a multiple of the number of bins? ### Part 1 — Bounds on the loss of one example For a single example, give a lower and an upper bound on its loss when the model predicts the correct class, and when it predicts an incorrect class. Say whether each bound can actually be reached. ```hint What the argmax forces If the true class has the largest of ten probabilities that sum to one, how small can its probability be? If some other class beats it, how large can it be? ``` #### What This Part Should Cover - The bounds for a correct prediction, with the argument - The bounds for an incorrect prediction, with the argument - Which bounds are attained, the role of ties, and what the overlap of the two ranges means ### Part 2 — Implement ACE ACE averages the gap between accuracy and predicted probability over bins that each hold the same number of predictions. For $N$ predictions, $K$ classes and $R$ bins: $$\mathrm{ACE} = \frac{1}{KR} \sum_{k=1}^{K} \sum_{r=1}^{R} \left| \mathrm{acc}(r, k) - \mathrm{conf}(r, k) \right|$$ For each class $k$, sort the $N$ predicted probabilities of class $k$ and cut them into $R$ consecutive bins of equal count. In bin $r$, $\mathrm{acc}(r, k)$ is the fraction of examples whose label is $k$, and $\mathrm{conf}(r, k)$ is the mean predicted probability of class $k$. Implement a function that takes the probability vectors, the labels and $R$, and returns ACE. ```hint Equal counts Decide how the sorted examples are cut into R groups when N is not a multiple of R, and what happens to examples with exactly equal probabilities at a cut. ``` #### What This Part Should Cover - Correct per-class sorting, equal-count binning and averaging - Handling of remainders, ties and invalid input - Time and memory cost ### Part 3 — Write test cases Write the test cases you would use to trust this implementation. ```hint Cases you can compute by hand Look for inputs whose ACE you know without running the code, and for properties that must hold for every input. ``` #### What This Part Should Cover - Small inputs with hand-computed expected values, including the extremes - Properties that must hold for any input - Edge cases and invalid inputs ### Part 4 — Why the ACE curve is spiky Plotted, the ACE results look spiky. Explain why, and how you would smooth the curve. Discuss what increasing the number of bins does, and what else you could do. #### Clarifying Questions for this Part - Is the curve the per-bin gap plotted against confidence, or ACE plotted against the number of bins or across training checkpoints? ```hint Count what is in a bin Estimate how many examples land in each bin for an accurate MNIST model, and what one example crossing a bin boundary does to the result. ``` #### What This Part Should Cover - The sources of noise and discontinuity in the estimate - The trade-off that the number of bins controls - Smoothing options and what each one costs ### What a Strong Answer Covers - Tight loss bounds derived from the argmax constraint, with attainability and ties handled - A correct, readable ACE implementation that matches the definition - Tests that combine hand-computed values, invariants and edge cases - A statistical explanation of the spikes, including the bias and variance of binned estimates - Practical smoothing choices and their trade-offs ### Follow-up Questions - How would equal-width bins, as in expected calibration error, behave on an accurate MNIST model compared with ACE's equal-count bins? - If you apply temperature scaling to the logits, which of your Part 1 bounds change, and what happens to ACE? - Why is the binned absolute gap biased upward even for a perfectly calibrated model, and how would you correct for it? - How would you report a confidence interval for ACE?