Using Ontologies To Improve Performance In Massively Multi-label Prediction Models

TL;DR

This paper introduces a novel neural network output layer that leverages ontologies to improve prediction accuracy for rare labels in massively multi-label classification tasks, such as disease and protein function prediction.

Contribution

It proposes a Bayesian network of sigmoids that incorporates ontology relationships to enhance learning for rare labels in multi-label models.

Findings

Significant improvements in AUROC for rare labels

Enhanced average precision for infrequent classes

Effective application to disease and protein function prediction

Abstract

Massively multi-label prediction/classification problems arise in environments like health-care or biology where very precise predictions are useful. One challenge with massively multi-label problems is that there is often a long-tailed frequency distribution for the labels, which results in few positive examples for the rare labels. We propose a solution to this problem by modifying the output layer of a neural network to create a Bayesian network of sigmoids which takes advantage of ontology relationships between the labels to help share information between the rare and the more common labels. We apply this method to the two massively multi-label tasks of disease prediction (ICD-9 codes) and protein function prediction (Gene Ontology terms) and obtain significant improvements in per-label AUROC and average precision for less common labels.