← Research

Hierarchical cross-entropy loss

When labels form a hierarchy, standard cross-entropy treats a prediction that is more specific than the label as simply wrong. This loss scores each class together with its subclasses, so any prediction inside the label's subtree counts.

Cell types are organised in an ontology, in which a CD4+ T cell is a T cell and a T cell is a lymphocyte. Single-cell studies annotate their cells at different depths of it: one study says T cell where another says CD4+ T cell. The standard cross-entropy treats every label as a separate class, so a model that predicts CD4+ T cell for a cell annotated as T cell is penalised as if it were wrong.

The hierarchical cross-entropy loss scores each cell type by its predicted probability plus the probabilities of all its descendants, and the loss is minus the log of the annotated type’s score. Probability flows up the ontology, and any prediction inside the annotated subtree counts, which is also how annotations are evaluated. In code it is one product with the ontology’s reachability matrix, so training costs the same.

Each node fills up to its score: its own probability plus that of its descendants. It opens on the paper's example, where T cell has probability 0.2 but, with its subtypes, a score of 0.8. Click a node to annotate the cell with it and compare the two losses on the same prediction.

We tested it the way atlases are used: train on the atlas as it was in May 2023, then annotate the studies added by December. On those new studies every model loses a large part of its accuracy. With the hierarchical loss they win back about half of it, with no change to the model and no extra computation.

The largest gains are on cell types inside the ontology, where related types were annotated in training, and the smallest on isolated leaves. So a new study helps a model most when it connects cell types that are already annotated, not only when it adds more cells.

The same idea applies wherever labels form a hierarchy. We later took it to image classification as Hierarchy-Aware Cross-Entropy, which also improves image classifiers on most of the architectures and datasets we tried.

Papers