Learning Structured Representations by Embedding Class Hierarchy with Fast Optimal Transport

TL;DR

This paper introduces a fast optimal transport-based method for embedding class hierarchies into feature representations, improving over previous approaches by better capturing multi-modal class distributions and achieving linear time complexity.

Contribution

The paper proposes an exact EMD-based regularizer within the CPCC framework and introduces the Fast FlowTree algorithm for efficient computation, advancing structured knowledge embedding.

Findings

The EMD-based method generalizes previous CPCC approaches.

Fast FlowTree algorithm achieves linear time complexity.

Method maintains competitive performance across datasets.

Abstract

To embed structured knowledge within labels into feature representations, prior work [Zeng et al., 2022] proposed to use the Cophenetic Correlation Coefficient (CPCC) as a regularizer during supervised learning. This regularizer calculates pairwise Euclidean distances of class means and aligns them with the corresponding shortest path distances derived from the label hierarchy tree. However, class means may not be good representatives of the class conditional distributions, especially when they are multi-mode in nature. To address this limitation, under the CPCC framework, we propose to use the Earth Mover's Distance (EMD) to measure the pairwise distances among classes in the feature space. We show that our exact EMD method generalizes previous work, and recovers the existing algorithm when class-conditional distributions are Gaussian. To further improve the computational efficiency of our method, we introduce the Optimal Transport-CPCC family by exploring four EMD approximation variants. Our most efficient OT-CPCC variant, the proposed Fast FlowTree algorithm, runs in linear time in the size of the dataset, while maintaining competitive performance across datasets and tasks. The code is available at https://github.com/uiuctml/OTCPCC.