Output Heads: Connecting a Network to a Task

A backbone produces a representation, and the output head converts that representation into values aligned with the task and loss.

Its shape and its semantics both must be explicit, because a single length-five vector could equally represent five class logits, five independent labels, five regression targets, or parameters of a distribution. Nothing in the tensor records which of those five things it is. That fact lives only in the loss, the target encoding, and the code that reads the output downstream.

A head can be asked to say much more than one number. One regression head emits an observation-noise term alongside the prediction. The loss then divides each squared error by a predicted variance, and pays a log-variance penalty for the privilege. Kendall and Gal, in 2017, called that “learned attenuation”: the model can declare a pixel intrinsically ambiguous and be charged less for missing it. They report that this “improve[s] model performance by 1 − 3% over non-Bayesian baselines”. They are equally clear about what the extra number does not cover. “Out-of-data examples, which can be identified with epistemic uncertainty, cannot be identified with aleatoric uncertainty alone.” The head can predict its own noise. It cannot predict what it has never seen.

The final transformation should reflect relationships among targets.

Scalar regression emits one unrestricted or constrained numeric value. Binary classification emits one logit, often transformed with sigmoid for interpretation. Multiclass classification emits one logit per mutually exclusive class, often normalized with softmax. Multilabel classification emits one independent logit per label, because several labels may be true at once. A distributional output emits several values that parameterize uncertainty, mixtures, intervals, or quantiles.

The list is not really five layer types. It is five different claims about how the coordinates of one output vector relate to each other. Do they compete? Are they independent? Do they name an answer at all, or only describe a spread? The sections below take the two claims that are easiest to get wrong: exclusivity, and the assumption that a number between 0 and 1 is a probability.

Both can produce a vector. The relationship among the coordinates is what changes the head and the loss.

A softmax head asserts that exactly one coordinate is correct. The field's most-used benchmark does not honour that assertion. Tsipras and colleagues measured it in 2020: “While each image is associated with a single label, we find that more than one fifth of ImageNet images contain objects from multiple classes.” On exactly those images, top-1 accuracy “drops by more than 10% across all models”. And for 650 of 2,156 multi-object images, the ImageNet label is not the main object a human annotator would name. A second group re-annotated the same validation set that year and put the figure at approximately 29% of images — images that either contain multiple objects or fall in a category matching multiple synonym labels. Their re-annotation assigned 57,553 labels to 46,837 images and discarded 3,163 images that received none. Two independent groups, a fifth to a third of the data. The exclusivity is a property of the head, not of the pictures.

Where co-occurrence is the point, the head is built for it. CheXpert, published in 2019, is 224,316 chest radiographs of 65,240 patients labelled for 14 co-occurring observations. The model emits a probability per observation rather than one normalized vector. A patient may have cardiomegaly and a pleural effusion at the same time, and no amount of softmax temperature will let a competing head say so. Because each coordinate is its own decision, each is judged against its own human operating point: three radiologists, individually, on a 500-study test set. Irvin and colleagues report where the model lands. “On Cardiomegaly, Edema, and Pleural Effusion, the model ROC and PR curves lie above all 3 radiologist operating points.” The MIMIC-CXR team, describing CheXpert for their own readers, records the same dataset independently: “A total of 224,316 chest radiographs for 65,240 patients admitted to Stanford Hospital were released with the CheXpert labeler … released with 14 labels”. They adopt that 14-label vocabulary for their own data too.

So the contrast is concrete. Multiclass: classes compete, softmax couples the probabilities, and the target is usually one class index. One document language is a fair example. Multilabel: the labels are not one exclusive choice, and sigmoid is applied per coordinate. Thresholds may differ from label to label, as they must when each of 14 observations has its own radiologist operating point. An image containing a car and a person is the ordinary case, not the awkward one.

Product logic often belongs after the network head. The network emits scores — logits, values, or distribution parameters. A transformation gives them range: sigmoid, softmax, decoding, or constraint functions. A calibration check tests whether that range means anything on the data the system will actually see. A policy then applies thresholds, costs, abstention, ranking, or human review. Monitoring measures the consequences of those decisions rather than only the raw model metric.

A regulator has written that separation into law. On 11 April 2018 the FDA granted De Novo request DEN180001 for IDx-DR, the first autonomous AI diagnostic authorised in the United States. The order created a new device type, “retinal diagnostic software device”, and attached special controls to it. One is a labelling requirement covering “The type of imaging data used, what the device outputs to the user, and whether the output is qualitative or quantitative.” That is the shape-and-semantics contract of the first section, made a condition of marketing. A second control requires clinical testing of sensitivity, specificity, PPV and NPV “across the range of available device outcomes” — the whole curve, not one convenient operating point. The pivotal trial behind the authorisation enrolled 900 subjects and cleared its pre-specified superiority endpoints: sensitivity 87.2% (95% CI 81.8-91.2%), specificity 90.7% (95% CI 88.3-92.7%), imageability 96.1% (95% CI 94.6-97.3%).

And here is what happens when a threshold travels without that check. Epic's proprietary sepsis model was validated externally in 2021 by Wong and colleagues, on 38,455 hospitalisations of 27,697 patients at Michigan Medicine between 6 December 2018 and 20 October 2019. At the vendor-suggested alert threshold of 6, the model gave sensitivity 33%, specificity 83% and PPV 12%, with a hospitalisation-level area under the ROC curve of 0.63 (95% CI 0.62-0.64). It missed 1,709 of 2,552 sepsis patients while “generating alerts for an ESM score of 6 or higher for 6971 of all 38 455 hospitalized patients (18%)”. The abstract concludes that “the ESM has poor discrimination and calibration in predicting the onset of sepsis”. Move the same head and the same number 6 to two county emergency departments and the behaviour changes again. A separate group measured sensitivity 14.7% and specificity 95.3% at that threshold in 2024, noting that “Alerting occurred at an Epic recommended model threshold of 6.” Thirty-three percent sensitivity in one hospital system, 14.7% in another. Identical score, identical cut-off. The head was constant; the decision it produced was not. That is precisely why steps three to five belong outside it.

Some tasks need structure or uncertainty that one scalar cannot express.

AlphaFold 2 shipped such a head into production and then published the evidence for it. The architecture description is one sentence in Nature: “Predictions of side-chain χ angles as well as the final, per-residue accuracy of the structure (pLDDT) are computed with small per-residue networks on the final activations at the end of the network.” The model does not merely emit a structure. It emits a confidence for every residue of it. That extra number is checked rather than asserted — the 2021 paper reports a least-squares fit of lDDT-Cα = 0.997 × pLDDT − 1.17, with Pearson's r = 0.76 over n = 10,795 protein chains. It also ships with an operating rule attached. EMBL-EBI, which co-operates the AlphaFold Protein Structure Database, tells users that “The predicted local distance difference test (pLDDT) is a per-residue measure of local confidence. It is scaled from 0 to 100”, and publishes the four bands to read it by: very high above 90, confident 90-70, low 70-50, very low below 50. A distributional head, a calibration check against true accuracy, and a decision policy, in one worked example.

Confirm target exclusivity, units, missing-target handling, feasible ranges, and the exact tensor shape, then test the head and loss together on a tiny batch with known targets.

Each of these has a documented failure behind it. Target exclusivity is the check ImageNet fails: more than one fifth of its images carry objects from several classes, while the head is built to name exactly one. Feasible ranges and declared semantics are the check the FDA turned into a marketing condition for IDx-DR, requiring the labelling to state what the device outputs and whether that output is qualitative or quantitative. Whether the loss is fed logits or probabilities is settled by reading the docstring. PyTorch's CrossEntropyLoss expects unnormalized logits, and passing it softmax output duplicates a transformation the loss already performs.

Also inspect whether constraints should be hard-coded or merely encouraged. The difference is not cosmetic. A soft penalty may still permit impossible outputs at inference.