Learning Structured Representations by Embedding Class Hierarchy with Fast Optimal Transport

## Abstract

Abstract To embed structured knowledge within labels into feature representations, prior work (Zeng et al., 2022 ) proposed to use the C o p henetic C orrelation C oefficient (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 O ptimal T ransport-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.

## 1 Introduction

Supervised classification problems have been a cornerstone in machine learning. While traditional classification methods treat class labels as independent entities, there has been a growing interest in leveraging semantic similarity among classes to improve model performance while following existing label dependency. This is particularly important in fine-grained classification tasks where labels are inherently living on a hierarchy, such as in the CIFAR (Krizhevsky and Hinton, ) and ImageNet (Deng et al., ) datasets. To capture the hierarchical relationship between classes, Zeng et al. () made a significant contribution by embedding the tree metric of a given class hierarchy into the feature space. They introduced a regularizer, the Cophenetic Correlation Coefficient (CPCC), to correlate the tree metric with the Euclidean distance in the class-conditioned representations. This approach of encoding hierarchical relationship in the latent space not only led to more interpretable features compared to other hierarchical classification methods (Shen et al., ), but also improved some generalization performance across multiple datasets.

Despite the value of this framework for embedding class hierarchies, it presents a notable limitation. Specifically, Zeng et al. () uses the $\ell_{2}$ norm between two class Euclidean centroids as the distance of two class distributions. However, this approach overlooks the difference of samples within the same class, and thus cannot fully capture the global structural information given pairs of class-conditional distributions, especially when they are multi-modal.
On the other hand, the use of Optimal Transport (OT) as an alternative to Euclidean centroids has first gained attention in the context of clustering, where Euclidean centroids neglect the intrinsic geometry of data distributions (Cuturi and Doucet, ). For the same motivation, we propose addressing the limitation of $\ell_{2}$-CPCC (Zeng et al., ), by integrating the OT distance, also known as the Wasserstein distance, and its variants, as a more effective metric for measuring the distance between class distributions in the problem of embedding hierarchical representations.

Our contributions are as follows: (i) We first generalize the $\ell_{2}$-CPCC with EMD, and propose a family of OT-CPCC algorithm to address the limitation of using $\ell_{2}$ centroids that misrepresents class relationships. We provide a detailed comparative analysis of $\ell_{2}$-CPCC versus OT-CPCC about their forward and backward computation. (ii) We propose a novel linear time OT approximation algorithm, FastFT, and apply it in the CPCC framework to learn hierarchical representations that is significantly more efficient than the cubic time exact EMD solution. (iii) We empirically analyze the advantage of OT-CPCC over $\ell_{2}$-CPCC across a wide range real-world datasets and tasks.

## 2 Preliminaries

#### Notation and Problem Setup

Throughout the paper we focus on the supervised $k$-class classification problem. We denote a sample of input $x$ with label $y$ living in the space of ${\mathcal{X}}\subseteq\mathbb{R}^{p}$ and ${\mathcal{Y}}=[k]\vcentcolon=\{1\ldots,k\}$. The dataset ${\mathcal{D}}$ consists of $N$ data pairs $\{(x_{i},y_{i})\}_{i=1}^{N}$.

Given the label space ${\mathcal{Y}}$, we have a corresponding label tree ${\mathcal{T}}$ that characterizes the relationship between different subsets of ${\mathcal{Y}}$. In particular, each node of ${\mathcal{T}}$ is associated with a subset of ${\mathcal{Y}}$: each leaf node $v$ is associated with a particular class in ${\mathcal{Y}}$; any internal node corresponds to the subset that is the union of its leaf children classes; the root node of ${\mathcal{T}}$ corresponds to the full ${\mathcal{Y}}$. The metric function of ${\mathcal{Y}}$ is called *tree metric* $t:{\mathcal{Y}}\times{\mathcal{Y}}\rightarrow\mathbb{R}_{+}$ which measures the similarity of two class labels $u,v$ based on the shortest path distance of two *leaf nodes* on ${\mathcal{T}}$. Intuitively, the tree ${\mathcal{T}}$ could be understood as a dendrogram of ${\mathcal{Y}}$ that characterizes the similarity of different classes.

Let ${\mathcal{Z}}\subseteq\mathbb{R}^{d}$ be the representation/feature space. A deep network $h$ is the composition of two functions $h_{w,\theta}=g_{w}\circ f_{\theta}:{\mathcal{X}}\rightarrow\Delta_{k}$ where $\Delta_{k}$ is a $k-1$ dimensional probability simplex. The feature encoder $f_{\theta}$ contains all layers until the penultimate layer of a neural network. On top of that, $g_{w}:{\mathcal{Z}}\rightarrow\Delta_{k}$ is a linear classifier with weights $w$. Let $q$ be the one-hot encoding vector of $y$. The parameters $\theta$ and $w$ are learned by minimizing ${\mathcal{L}}$ combining the cross entropy loss $\ell_{\text{CE}}$ with an optional regularization term ${\mathcal{R}}$:

$$ ${\mathcal{L}}(x,y;w,\theta)=\sum\limits_{i=1}^{k}-q_{i}\log h_{w,\theta}(x)_{i }+\lambda{\mathcal{R}}(x,y,\theta).$ (1) $$

The hyperparameter $\lambda$ is the regularization factor which controls the strength of ${\mathcal{R}}$. The goal of the problem is to embed the class hierarchical relationship in the feature space. At a high level, the distance of class-conditioned representations should approximate the tree metric of classes in ${\mathcal{T}}$. This extra constraint can be quantitatively characterized by CPCC as the regularizer term ${\mathcal{R}}$ that we will introduce next.

### 2.1 Cophenetic Correlation Coefficient

The Cophenetic Correlation Coefficient (CPCC) (Sokal and Rohlf, ) was introduced as a regularization term in Zeng et al. () to embed the class hierarchy ${\mathcal{T}}$ in the feature space. CPCC has the same formulation as Pearson’s correlation coefficient that measures the linear dependence between two metrics $t$ and $\rho$:

$$ $\text{CPCC}(t,\rho)=\frac{\sum\limits_{u<v}(t(u,v)-\bar{t})(\rho(u,v)-\bar{ \rho})}{(\sum\limits_{u<v}(t(u,v)-\bar{t})^{2})^{-1/2}(\sum\limits_{u<v}(\rho( u,v)-\bar{\rho})^{2})^{-1/2}},$ (2) $$

where $\rho$ is a distance function between $u,v$ on feature space ${\mathcal{Z}}$, and $\bar{t},\bar{\rho}$ are the average of each metric. Let ${\mathcal{D}}_{u}=\{(x^{u}_{i},u)\}_{i=1}^{m},{\mathcal{D}}_{v}=\{(x^{v}_{i},v
)\}_{i=1}^{n}\subseteq{\mathcal{D}}$ be datasets with the same class label $u,v$. When $\ell_{2}$-CPCC is used as a regularizer in Eq. 1,

$$ $\rho(u,v)\vcentcolon=\left\|\frac{1}{m}\sum_{(x^{u}_{i},u)\in{\mathcal{D}}_{u} }f(x^{u}_{i})-\frac{1}{n}\sum_{(x^{v}_{j},v)\in{\mathcal{D}}_{v}}f(x^{v}_{j}) \right\|_{2},$ $$

which is the $\ell_{2}$ distance of the mean of the two class-conditional features. We call $z_{i}^{u}$ as the $i$-th feature vector with label $u$, when the context is clear, we drop the superscript $u$ for convenience. By maximizing CPCC ($\lambda<0$), $\rho$ follows the relation in $t$: classes sharing a close common ancestor will be grouped together while classes with a higher common ancestor are pulled away with respect to Euclidean distance.

### 2.2 Exact Earth Mover’s Distance

Earth Mover’s Distance (EMD), or Optimal Transport distances, aims to find the optimal transportation plan between two discrete measures. For ${\mathcal{D}}_{u}$, let $a=(a_{1},\dots,a_{m})\in\Delta_{m}$ which is not necessarily uniform. Let $b$ defined similarly for ${\mathcal{D}}_{v}$, and $\delta_{z}$ is the Dirac delta at $z$. Define two discrete measures as $\mu=\sum_{i=1}^{m}a_{i}\delta_{z_{i}},\nu=\sum_{j=1}^{n}b_{j}\delta_{z_{j}}$. Then, EMD between probability distributions defined over ${\mathcal{Z}}$, is the objective value of the following linear program:

$$ $\min\limits_{{\bm{P}}}\ \langle{\bm{P}},{\bm{D}}\rangle,\quad\mathrm{\,s.t.\ } {\bm{P}}{\bm{1}}_{m}=a,~{}{\bm{P}}^{\top}{\bm{1}}_{n}=b.$ (3) $$

In Eq.3, $\langle\cdot,\cdot\rangle$ is the Frobenius inner product, ${\bm{1}}_{N}$ is the $N$-dimensional all ones vector, ${\bm{D}}$ is the $m\times n$ pairwise distance matrix where ${\bm{D}}_{ij}=\left\lVert z_{i}-z_{j}\right\rVert$, and ${\bm{P}}$ is called the flow matrix in the computer science literature, or optimal transport plan in the OT literature, the same size with ${\bm{D}}$ with nonnegative entries. Intuitively, this is a minimum cost flow problem. Both source and target flow sum to $1$, and we want to find the optimal transportation plan ${\bm{P}}$ with the minimum cost with respect to ground similarity ${\bm{D}}$. The norm $\lVert\cdot\rVert$ depends on the ground metric of ${\bm{D}}$. We use the $\lVert\cdot\rVert_{2}$ norm as the ground metric throughout this paper.

When the dimension of the feature is 1, fast computation of the EMD distance exists (Houry et al., ). Instead of solving a linear program in Eq.3, when $m=n$, $a_{i}=b_{i}=1/n$ for all $i$, the closed form solution for 1d EMD is: $\text{EMD}_{\text{1d}}(\mu,\nu)=\frac{1}{n}\sum_{i=1}^{n}\|z_{(i)}-z^{\prime}_
{(i)}\|$, and $z_{(i)}$ is the $i$-th order statistic of the sequence $\{z_{1},\cdots z_{n}\}$ after sorting. When $m\neq n$, $\text{EMD}_{\text{1d}}$ generalizes to solving inverse CDF of $\mu,\nu$ (Peyré and Cuturi, ) which requires approximation. To avoid this approximation, *after sorting*, we can compute the flow matrix ${\bm{P}}$ (Hoffman, ) efficiently via a greedy algorithm in linear time (Alg.1) and use the inner product objective of EMD in Eq.3 to compute the result. This greedy flow matching algorithm that generates the optimal transport plan can be used for any $m,n$, including the special case when $m=n$, and fit any non-uniform weights of $a,b$.

Figure: Algorithm 1 Greedy Flow Matching

## 3 CPCC with Optimal Transport

### 3.1 OT-CPCC Family

Despite its simplicity, *the use of class centroids $\ell_{2}$ distance as a measure of class distances has two limitations.* First, since there is no guarantee of the distribution of class conditional features, the class mean difference may not accurately reflect the average pairwise Euclidean distance between all the pairs of sample. For example, in Fig.1, class $v$ is a non-unimodal (bi-modal in this example) distribution, which is very likely to happen when class $v$ is an abstract concept, such as coarse classes like household furniture in CIFAR100 (Krizhevsky and Hinton, ) or the collection of diverse dog breeds in ImageNet Deng et al. (). The blue $\ell_{2}$ line can be a biased estimation of the average of red lines. We will see later (c.f. Sec.3.3) that the direction of blue vector plays an important role in the feature geometry in $\ell_{2}$-CPCC, and therefore we misrepresent the class distribution in $\ell_{2}$-CPCC. Second, it is known (Linial et al., ) that, for certain trees such as a basic star graph with three edges, it is impossible to embed tree metric $t$ into $\ell_{2}$ exactly, no matter how large feature dimension $d$ is. It implies if we use the Euclidean *centroids* to represent a class, we can never reach the optimal CPCC value $1$ during training, which always cause some loss of hierarchical information in the learning of structured representation.

Figure: Figure 1: Comparison between EMD (weighted sum of red lines) and the class mean $\ell_{2}$ distance (blue line).
Refer to caption: x1.png

How to overcome the above limitations? Instead of finding a one-to-one correspondence of a label and a point estimate of a class, we represent a label class with all *samples* that belong to the class, which leads to EMD-CPCC.
We call the set of data-weight pairs containing $m$ class-$u$ conditioned features as $\mu=\{(z_{i},a_{i})\}_{i=1}^{m}$ and define $\nu$ similarly for the dataset containing $n$ class-$v$ conditioned features. The weight $a_{i}$ can be set as uniform ($1/m$) by default, or non-uniform if external prior knowledge is available. We propose EMD-CPCC by implementing $\rho(u,v)\vcentcolon=\text{EMD}(\mu,\nu)$. In Eq.3, ${\bm{P}}$ depends on ${\bm{D}}$ derived from each *sample* in two datasets. Besides, EMD is actually a generalization of $\ell_{2}$ method:

###### Proposition 3.1 (EMD reduces to $\ell_{2}$ for Gaussians with same covariance) .

$\mathrm{EMD}({\mathcal{N}}(\mu_{z},\Sigma),{\mathcal{N}}(\mu_{z^{\prime}},
\Sigma))=\|\mu_{z}-\mu_{z^{\prime}}\|$.

Although EMD has many appealing properties, the computational cost of EMD is high, where the time complexity of linear programming is cubic time with respect to the data size. Many following work provides approximation methods to reduce the computational cost, such as Sinkhorn (Cuturi, ), Sliced Wasserstein Distance (SWD) (Bonneel et al., ), Tree Wasserstein Distance (TWD) (Yamada et al., ; Le et al., ; Takezawa et al., ), and FlowTree (FT) (Backurs et al., ) (detailed discussion of approximation methods in App.A). Similarly, we can replace $\rho$ with all other EMD approximation methods, and we call all of them as the OT-CPCC family. As a side remark, since all approximation methods provide a ${\bm{P}}$ satisfying constraints in Eq.3 and since EMD is lower bounded by $\ell_{2}$ (see proof of Prop.3.1 in App.B), all OT-CPCC methods are still lower bounded by $\ell_{2}$.

###### Proposition 3.1 (EMD reduces to ℓ 2 subscript ℓ 2 \ell_{2} roman_ℓ start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT for Gaussians with same covariance) .

### 3.2 Fast FlowTree

#### FlowTree

(Backurs et al., ) is a tree-based method where the structure of the tree ${\mathcal{T}}$ is learned using QuadTree (Samet, ) to approximate the Euclidean space. Once the tree has been built, it firsts compute the optimal flow matrix ${\bm{P}}_{t}$ using $t$ as the ground metric with a linear time bottom-up as shown in Alg.2 (Kalantari and Kalantari, ) (see App.Fig.8 in for visualization). Then, it estimates EMD with $\langle{\bm{P}}_{t},{\bm{D}}_{\text{$\ell_{2}$}}\rangle$. Given the tree, the time complexity of computing the FT is $O((m+n)d\log(d\Phi))$ where $\Phi$ is the maximum Euclidean distance of any two feature vectors.

#### Our Fast FlowTree

The main computational bottleneck of FT is the construction of the tree. However, in our case, we have access to ${\mathcal{T}}$ given as prior knowledge, and thus we can skip the tree construction step. To apply tree-based EMD approximation methods, instead of building a QuadTree in FT, we first construct an *augmented tree* by extending the leaves of the original label tree ${\mathcal{T}}$ with samples. Specifically, since each leaf node in ${\mathcal{T}}$ corresponds to a class, we extend each leaf node by adding all samples from the corresponding class, as demonstrated in Figure 2. We set all extended edges to have edge weight $1$, and assign the instance weight $a_{i}$ for each data point. The resulting algorithm that applies Alg.2 to the *augmented* label tree is called Fast FlowTree (FastFT).

One key insight that allows FastFT to be more efficient than FT is that the tree metric between any two data samples from different classes is always identical in the augmented tree ${\mathcal{T}}$, which allows us to reduce the computation of ${\bm{P}}_{t}$ to 1d OT greedy flow matching using Alg.1. The time complexity of FastFT is $O((m+n)d)$, which is optimal among all OT variants.

###### Theorem 3.2 (Correctness of Fast FlowTree) .

For any augmented label tree ${\mathcal{T}}$, Alg. 1 and Alg. 2 return the same EMD approximation. The Fast FlowTree can be computed in $O((m+n)d)$.

Please refer to App.C for proof details of Theorem 3.2. Note that all OT-CPCC methods inherently need $\Omega((m+n)d)$ because we access each of $m+n$ data point once and look at each entry in $d$-dimensions. Additionally, FastFT can be applied to any label tree with no restriction on structure or edge weights. To control the importance of data points, the flexibility of FastFT can be achieved by tuning $a_{i}$.

Figure: Table 1: Time complexity comparison of different CPCC methods. Let $b$ be the batch size, $I$ be the number of iterations for iterative methods, $p$ be the number of projections for SWD, $\Phi$ be the maximum Euclidean distance of any two feature vectors. We ignore the factor of $\min(b^{2},k^{2})$ shared for all methods due to batch computation. For simplicity we assume $m=n$. In the batch learning setting, $d$ can be much larger than $n$ particularly for the fine-grained classification problems.
Refer to caption: x2.png

###### Theorem 3.2 (Correctness of Fast FlowTree) .

#### Comparison of OT-methods on Synthetic Dataset

To further understand the behavior of different OT approximation methods, we present experiments on synthetic datasets in Fig.3 (see App.H for more details). We first measured the computation time for different OT methods across datasets of varying sizes. Second, we evaluated the error of different OT approximation methods, along with the $\ell_{2}$ distance of the mean of two datasets, in approximating EMD. Two different scenarios were considered, either both distributions were Gaussian or non Gaussian. Under Gaussian distribution, the OT methods’ results were closely aligned with the $\ell_{2}$ method, except for SWD, which supports Prop.3.1 and implies the similarity in performance between OT and $\ell_{2}$-CPCC under Gaussian distributions. In non-Gaussian distribution scenario, notable differences emerged between OT methods and $\ell_{2}$, where SWD is closer to $\ell_{2}$ measurement. FastFT demonstrated exceptional advantage as an approximation method, outperforming Sinkhorn in both computational speed and precision.

Figure: Figure 3: Efficiency and approximation error comparison of OT methods on synthetic datasets.
Refer to caption: x3.png

### 3.3 Differentiability Analysis of OT-CPCC

Since we use some non-differentiable algorithms during training (ex., linear programming in EMD, sorting in SWD, tree construction in TWD, greedy algorithm in FT and FastFT), it is unclear how does the backpropagation algorithm work exactly in this problem. Thanks to the Danksin’s theorem (Danskin, ) and the structure of the underlying optimization problem, we can still efficiently compute the gradients of OT-CPCC. For simplicity, let us consider the case of EMD-CPCC, i.e., $\rho=\text{EMD}$ exactly. The other approximate variants can be analyzed in a similar way.

The underlying challenge of the problem is to take derivatives of the optimal value of a minimization problem. We state the differentiability section of Danskin’s theorem formally and prove a lemma:

###### Theorem 3.3 ( Danskin ( 1966 ) ) .

Let $V(x)=\max_{z\in\mathcal{C}}f(x,z)$, and $Z_{0}(x)$ is the set of $z$ that achieve the maximum of $V$ at $x$. If $f:\mathbb{R}^{n}\times\mathcal{C}\rightarrow\mathbb{R}$ is continuous and $\mathcal{C}$ is compact, then the subgradient of $V$ is the convex hull of set $\{\frac{\partial f(x,z)}{\partial x}:z\in Z_{0}\}$. Particularly, if $Z_{0}=\{z^{*}\}$, $\frac{\partial V}{\partial x}=\frac{\partial f(x,z^{*})}{\partial x}$.

###### Lemma 3.4 .

Denote ${\bm{P}}^{*}$ as the optimal flow of EMD problem, then $\frac{\partial\rho_{\text{EMD}}}{\partial\theta}={\bm{P}}^{*}\cdot\frac{
\partial D}{\partial\theta}$.

###### Proof.

We rewrite the objective function as a maximization problem: $\rho_{\text{EMD}}=-\max_{{\bm{P}}\in\mathcal{C}}-\langle{\bm{P}},{\bm{D}}\rangle$.
Roughly speaking, Danksin’s theorem states that the derivative of maximization problems can be evaluated at the optimum. We can get (sub-)differential of $\rho_{\text{EMD}}$ with respect to $\theta$ with one step of chain rule:

$$ $\frac{\partial\rho_{\text{EMD}}}{\partial\theta}=\frac{\partial\rho_{\text{EMD }}}{\partial{\bm{D}}}\cdot\frac{\partial{\bm{D}}}{\partial\theta}=-\frac{ \partial{-\langle{\bm{P}}^{*},{\bm{D}}\rangle}}{\partial{\bm{D}}}\cdot\frac{ \partial{\bm{D}}}{\partial\theta}={\bm{P}}^{*}\cdot\frac{\partial{\bm{D}}}{ \partial\theta},$ $$

where the last term comes from the original flow of backpropagation containing the gradients with respect to network parameters. The key takeaway is the gradient of $\rho_{\text{EMD}}$ w.r.t. the network parameters does not require the gradient from ${\bm{P}}$ in the EMD setting, although ${\bm{P}}$ depends on $\theta$.
∎

###### Theorem 3.3 ( Danskin ( 1966 ) ) .

###### Lemma 3.4 .

###### Proof.

#### Differentiability for Other OT-CPCC Methods

The application of Danskin’s relies on the optimality of ${\bm{P}}^{*}$. How about the other approximation methods? The differentiability is not a problem in Sinkhorn as the iterative algorithm only involves matrix multiplication. But for the SWD, TWD, FT, the approximate solutions (the flow matrix ${\bm{P}}$) for these methods are not explicitly formulated as solutions to an optimization problem. Following EMD, except for the Sinkhorn variant, our method is to treat flow matrix ${\bm{P}}$ always as a constant for these approximation methods, and stop the gradient through ${\bm{P}}$. For example, $\frac{\partial\rho_{\text{FT}}}{\partial\theta}\vcentcolon={\bm{P}}_{\text{FT}
}\cdot\frac{\partial D}{\partial\theta}$.

To provide a more fine-grained understanding of the proposed method, after showing and defining the differentiability of OT-CPCC, in what follows we analyze the gradients of the proposed CPCC variants and compare them with the $\ell_{2}$-CPCC.

#### Gradient of ℓ 2 subscript ℓ 2 \ell_{2} roman_ℓ start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT -CPCC

If we want to know the difference between $\ell_{2}$-CPCC and EMD-CPCC, assume the same input and the same network parameters, ignoring the cross-entropy loss, the only difference of the gradient starts from the distance calculation between two classes.
Assume representations from two classes are $Z\in\mathbb{R}^{m\times d},Z^{\prime}\in\mathbb{R}^{n\times d}$.

$$ $\displaystyle\frac{\partial{\rho_{\text{$l_{2}$}}}}{\partial{Z}}$ $\displaystyle=\frac{\partial}{\partial{Z}}\left\lVert\mu_{z}-\mu_{z^{\prime}} \right\rVert=\frac{\mathbf{1}(\frac{1}{m}\mathbf{1}^{\top}Z-\frac{1}{n}\mathbf {1}^{\top}{Z^{\prime}})}{m\left\lVert\frac{1}{m}Z^{\top}\mathbf{1}-\frac{1}{n} {Z^{\prime}}^{\top}\mathbf{1}\right\rVert}.$ $$

Let $i$-th row of $Z$ be $Z_{i\cdot}$. Each row of $\frac{\partial{\rho_{\text{$l_{2}$}}}}{\partial{Z}}$ is the same, so

$$ $\frac{\partial{\rho_{\text{$l_{2}$}}}}{\partial{Z_{i\cdot}}}=\frac{\mu_{z}-\mu _{z^{\prime}}}{m\left\lVert\mu_{z}-\mu_{z^{\prime}}\right\rVert}.$ $$

#### Gradient of EMD-CPCC

Let ${\bm{D}}\in\mathbb{R}^{m\times n}$ be the cost matrix and ${\bm{P}}$ is any constant flow matrix:

$$ $\rho_{\text{EMD}}=\langle{\bm{P}}^{*},{\bm{D}}(Z,Z^{\prime})\rangle=\sum \limits_{i,j}{\bm{P}}^{*}_{ij}\left\lVert Z_{i\cdot}-Z^{\prime}_{j\cdot}\right\rVert.$ $$

Let’s calculate the derivative by rows of $Z$. We only keep the terms related to $Z_{i\cdot}$.

$$ $\displaystyle\frac{\partial{\rho_{\text{EMD}}}}{\partial{Z_{i\cdot}}}$ $\displaystyle=\frac{\partial}{\partial{Z_{i\cdot}}}\sum\limits_{j}{\bm{P}}_{ij }^{*}\left\lVert Z_{i\cdot}-Z_{j\cdot}^{\prime}\right\rVert=\sum\limits_{j=1}^ {n}{\bm{P}}_{ij}^{*}\frac{Z_{i\cdot}-Z_{j\cdot}^{\prime}}{\left\lVert Z_{i \cdot}-Z_{j\cdot}^{\prime}\right\rVert}.$ $$

Note that since ${\bm{P}}^{*}$ is the optimal solution of a linear program, it is a sparse matrix with at most $m+n-1$ non-zero entries (Peyré et al., , Proposition 3.4). As a comparison to the gradient of $\ell_{2}$-CPCC, which assigns equal importance to all data points and only depends on the direction of class centroids, the gradient of EMD-CPCC assigns different importance to each pair of data points and depends on the pairwise transport probability between data points. From this analysis, we can see that the gradient provided by EMD-CPCC is more fine-grained and informative than that of $\ell_{2}$-CPCC.

## 4 Experiments

For a fair comparison, we closely aligned our implementation with Zeng et al. () to ensure consistency. This section will cover dataset benchmarks, the multi-modality of features, downstream hierarchical classification retrieval, and interpretability analysis. A detailed description of training hyperparameters is provided in App.E, along with additional experiments, including the effect of label hierarchy structure, label hierarchy tree weights, optimal transport flow settings (distribution of $a$ in Eq.3), and comparisons with non-CPCC hierarchical methods, in App.G. Throughout this section, the Flat objective is the standard cross entropy loss in Eq.1 with $\lambda=0$.

**Table 2: Dataset Statistics.**
| Dataset | # train | # test |
| --- | --- | --- |
| CIFAR10 | 50K | 10K |
| CIFAR100 | 50K | 10K |
| INAT | 500K | 100K |
| L17 | 88K | 3.4K |
| N26 | 132K | 5.2K |
| E13 | 334K | 13K |
| E30 | 307K | 12K |

### 4.1 Datasets

We conducted experiments across $7$ diverse datasets, CIFAR10, CIFAR100 (Krizhevsky and Hinton, ), BREEDS benchmarks (Santurkar et al., ) including four settings LIVING17 (L17), ENTITY13 (E13), ENTITY30 (E30), NONLIVING26 (N26), and iNaturalist-mini (Van Horn et al., ) (INAT) to thoroughly evaluate the efficacy and influence of OT-CPCC methods. The dataset statistics is shown in Tab.3. Note that BREEDS and INAT are larger scale datasets with high resolution images. As subsets of ImageNet, BREEDS inherit a more complicated label hierarchy and more realistic, possibly multi-modal class distributions. The major difference from ImageNet is that BREEDS’s class hierarchy is manually calibrated so that it reflects the visual hierarchy instead of the WordNet semantic hierarchy in the original ImageNet, which removes ill-defined hierarchical relationships unrelated to vision tasks.

### 4.2 Existence of Multi-mode Class-conditioned Feature Distribution

Figure: Figure 5: The average best number of components of learned features (left) across training epochs on CIFAR100 for different CPCC methods (right) for the $\ell_{2}$-CPCC final trained features for all datasets.
Refer to caption: x6.png

To verify our hypothesis in Fig.1, we aim to determine whether class-conditioned clusters exhibit a multi-modal structure in practice. We track the class-conditioned features during training: at each epoch, for each class-conditioned feature cluster, we fit a Gaussian Mixture Model (GMM) with the number of components ranging from 1 to 10. We report the optimal number of components based on the lowest Akaike Information Criterion (AIC), averaged across all fine-level class labels. *If the best number of Gaussian components exceeds 1, this may indicate the presence of a multi-modal structure in the feature distributions*, assuming a sufficiently large batch size. From Fig.6, with default training hyperparameters in App.E, we observe that this phenomenon occurs across most datasets for both $\ell_{2}$-CPCC and OT-CPCC training, suggesting that using optimal transport distances in CPCC computation is more appropriate.

### 4.3 Hierarchical Classification and Retrieval

**Table 4: Generalization performance on coarse level. The prefix of metric represents the split of test set: s = source, t = target. The average number of three seeds are reported. See the detailed version including standard deviation in Tab.14,15 and performance on each BREEDS setting is in Tab.16,17. Due to the computation complexity, we do not report EMD results on training BREEDS from scratch.**
| Dataset | Objective | sAcc | tAcc | sMAP | tMAP | Dataset | sAcc | tAcc | sMAP | tMAP |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| CIFAR10 | Flat | 99.58 | 87.30 | 99.22 | 89.66 | INAT | 94.63 | 38.81 | 70.41 | 34.00 |
| FastFT | 99.61 | 87.79 | 99.91 | 93.04 | 94.66 | 39.43 | 72.90 | 35.63 |  |  |
| EMD | 99.61 | 87.45 | 99.88 | 93.07 | 94.64 | 41.01 | 73.87 | 35.58 |  |  |
| Sinkhorn | 99.61 | 87.41 | 99.87 | 92.60 | 94.30 | 36.94 | 68.75 | 34.92 |  |  |
| SWD | 99.56 | 87.63 | 99.36 | 90.12 | 94.52 | 39.38 | 75.13 | 38.11 |  |  |
| CIFAR100 | Flat | 86.17 | 43.09 | 82.07 | 42.09 | BREEDS | 93.85 | 82.43 | 74.47 | 66.50 |
| FastFT | 86.96 | 44.29 | 91.71 | 42.20 | 94.25 | 82.89 | 92.84 | 79.69 |  |  |
| EMD | 86.93 | 44.03 | 91.82 | 42.32 | - | - | - | - |  |  |
| Sinkhorn | 86.93 | 43.99 | 91.61 | 42.45 | 94.42 | 83.51 | 91.84 | 79.33 |  |  |
| SWD | 86.42 | 43.35 | 87.24 | 42.39 | 94.41 | 82.96 | 83.58 | 72.38 |  |  |

**Table 5: Generalization performance on fine level.**
| Dataset | Objective | sAcc | tAcc | sMAP | Dataset | sAcc | tAcc | sMAP |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| CIFAR10 | $\ell_{2}$ | 96.96 | 55.71 | 99.22 | INAT | 88.62 | 26.78 | 56.10 |
| FastFT | 96.90 | 55.99 | 99.24 | 88.49 | 27.10 | 56.21 |  |  |
| EMD | 97.05 | 56.12 | 99.24 | 88.68 | 26.78 | 56.83 |  |  |
| Sinkhorn | 96.95 | 54.89 | 99.27 | 88.08 | 26.77 | 51.56 |  |  |
| SWD | 96.96 | 59.21 | 99.45 | 88.46 | 26.78 | 54.19 |  |  |
| CIFAR100 | $\ell_{2}$ | 80.98 | 23.76 | 77.52 | BREEDS | 82.66 | 45.95 | 62.86 |
| FastFT | 81.09 | 24.71 | 78.72 | 82.58 | 45.75 | 63.95 |  |  |
| EMD | 81.32 | 23.15 | 78.88 | - | - | - |  |  |
| Sinkhorn | 80.99 | 23.53 | 78.51 | 82.99 | 46.87 | 61.54 |  |  |
| SWD | 80.50 | 26.18 | 86.82 | 82.90 | 46.33 | 76.24 |  |  |

#### Hierarchical Dataset Split

For downstream classification and retrieval, we split each dataset into source and target subsets, with both further divided into train and test sets, as illustrated in Fig.4. The source and target datasets share the same level of coarse labels but differ in fine labels, where the fine level refers to the more granular class labels that are exactly one level below the coarse labels. For example, in CIFAR10, the source set includes the fine labels (deer, dog, frog, horse) and (ship, truck), while the target set includes (bird, cat) and (airplane, automobile). Both sets share the same coarse labels: animal and transportation. We include details of hierarchical construction and splits for other datasets in App.D.

Given the hierarchical splits, we define two evaluation protocols. First, we pretrain on the source-train set and evaluate on the source-test set, where the train and test distributions are i.i.d. Second, we pretrain on the source-train set and evaluate on the target-test set, introducing a subpopulation shift (Santurkar et al., ) between the train and test distributions. For each protocol, we conduct classification and retrieval experiments at two levels of granularity.

#### Metrics

For classification experiments, we measure performance using accuracy (Acc). At the fine level, source accuracy is computed directly using any pretrained model, while coarse source accuracy is calculated by summing the fine-level probabilities for predictions that share the same coarse parent label. Since the source and target sets share the same coarse labels, coarse target accuracy can be similarly computed using the model pretrained on the source split. For fine target accuracy, after pretraining on the source set, we train a linear classifier on top of the frozen pretrained encoder for one-shot learning on the target-train set and evaluate it on the target-test set. We include the visualization of all types of classification tasks evaluation in Fig.9 in the Appendix.

For retrieval experiments, we use mean average precision (MAP) as the evaluation metric. Given a model pretrained on the source set, we compute the fine or coarse class prototypes by calculating the Euclidean centroid of the class-conditioned features. Using these class prototypes, we compute the cosine similarity score between the prototypes and test points, and evaluate whether test points are most similar to the class prototype with the same label via average precision. Since the source and target sets have different fine labels, we do not include target MAP in Tab.5.

#### Discussion

OT-CPCC outperforms $\ell_{2}$-CPCC consistently on fine level tasks: the oracle for the fine level tasks should be Flat, and Zeng et al. () reported that $\ell_{2}$-CPCC does not perform well for several fine level tasks due to the constraint on coarse level. Tab.5 shows that OT-CPCC methods shows better fine-level performance than $\ell_{2}$, while maintaining hierarchical information simultaneously. Besides, OT-CPCC outperforms the Flat baseline consistently on coarse level tasks as expected in Tab.4, which implies OT-CPCC methods successfully include the additional hierarchical information by optimizing the regularizer. In Fig.6, we visualize the learned hierarchical embeddings using UMAP (McInnes et al., ), with Euclidean distance as the similarity metric. We observe that $\ell_{2}$-CPCC displays a clearer coarse-level grouping pattern, likely due to the direct optimization of CPCC in Euclidean space. However, EMD-CPCC also retains a significant amount of hierarchical information and demonstrates better separation between fine-level classes, which aligns with the results observed in our classification and retrieval tasks.

### 4.4 Interpretability Evaluation

Figure: Figure 7: Test CPCC for various batch size, regularization strength, and architecture, from left to right. FastFT consistently performs the best across all settings.
Refer to caption: x7.png

A model is more interpretable when its representations accurately capture hierarchical information. CPCC is a natural measure of the precision of the hierarchical structure embedded in corresponding metric spaces. Unlike rank-based metrics, where only the relative order of distances matters, CPCC not only requires the correct ranking of distances (i.e., fine classes sharing the same coarse class should be closer than fine classes with different coarse parents), but is also influenced by the absolute values of the node-to-node distances, because CPCC is affected by the given label tree weights (see App.G.2 for more discussion). Therefore, we report CPCC, a stricter metric, on the test set. In Tab.3, we present $\ell_{2}$-CPCC for the Flat and $\ell_{2}$ objectives, and use the corresponding OT distances for the OT-CPCC experiments.

We observe that OT-CPCC preserves more hierarchical information, leading to better interpretability. As shown in Tab.3, in 5 out of 7 datasets, an OT-CPCC method achieves the best TestCPCC score. For E13 and E30, OT-CPCC methods slightly underperform compared to $\ell_{2}$, likely due to the BREEDS representations being less multi-modal than those of other datasets, as indicated in Tab.6.

#### Ablation Study

In our experiments with CIFAR100, shown in Fig.7, we investigated the impact of batch size, the regularization factor $\lambda$ in Eq.1, and architecture on TestCPCC. Given CIFAR100’s large number of fine classes, the average sample size per class is relatively small when applying CPCC with a standard batch size of 128. Consequently, there is minimal difference between using the centroid of a few data points, as in $\ell_{2}$-CPCC, and using pairwise distances for each data point in OT-CPCC. Therefore, larger batch sizes allow OT-CPCC to better capture class distributions, potentially leading to a more significant advantage in the quality of the learned hierarchy. We compared TestCPCC across various settings, including batch sizes of [64, 128, 512, 1024], $\lambda$ values of [0.01, 0.05, 0.1, 0.5, 1, 2], and backbones such as ResNet18, ResNet34, ResNet50 (He et al., ), ViT-B/16, and ViT-L/16 (Dosovitskiy, ), using $\ell_{2}$, EMD, or FastFT CPCC. Across all comparisons, we observe that FastFT > EMD > $\ell_{2}$ in terms of the quality of the learned hierarchy, highlighting the advantage of optimal transport-based methods.

## 5 Related Work

#### Learning Hierarchical Representations

We discuss methods that operate on feature representations, thereby directly influencing the geometry of the representation space when hierarchical information is provided. DeViSE (Frome et al., ) is an early multimodal work to align image and text embedding evaluated on a hierarchical precision metric to reflect semantic knowledge in the representation. Ge et al. (); Zhang et al. () modify triplet loss to capture class similarity for metric learning or image retrieval. Barz and Denzler (); Garg et al. () embed classes on hypersphere based on hierarchical similarity, and similar ideas have been explored in the hyperbolic space (Kim et al., ) where tree can be embedded with low distortion. Yang et al. () provides a multi-task proxy loss for deep metric learning. Nolasco and Stowell () proposes a ranking loss to make Euclidean distance of feature close to a target $d$-dim sequence determined by the rank of tree metric. We work on the extension of Zeng et al. () which is a special type of hierarchical methods where a regularizer is used. For a broader discussion on various hierarchical methods, we direct readers to the summary provided by Bertinetto et al. ().

#### OT Methods and Hierarchical Learning

The application of Optimal Transport (OT) in hierarchical learning contexts is explored in a limited number of studies. Ge et al. (), for instance, substitutes the conventional cross-entropy with a novel loss function that optimally transports the predictive distribution to the ground-truth one-hot vector, modifying the ground metric as tree metric and thus utilizing the hierarchical structure of the classes. Additionally, SEAL (Tan et al., ) employs the TWD to construct a latent hierarchy, enhancing traditional supervised and semi-supervised learning methodologies. It can be demonstrated that, given a tree ${\mathcal{T}}$, the optimization process for TWD aligns with $\ell_{1}$ loss optimization and is upper bounded by the SumLoss baseline introduced by Zeng et al. (). Consequently, these methodologies can be categorized as hierarchical loss functions rather than regularization methods discussed in this work.

## 6 Conclusion

Building upon Zeng et al. () using CPCC as a regularizer to learn structured representations, we have identified and targeted the limitations inherent in $\ell_{2}$-CPCC. Our proposed OT-CPCC family provides a more nuanced measurement of pairwise class distances in feature space, effectively capturing complex distributions theoretically and empirically. Furthermore, our novel linear FastFT algorithm itself is interesting in nature for OT communities. FastFT-CPCC can be viewed as a successful application in learning tasks, and we welcome future work to investigate more on the properties of our greedy approach.

## Acknowledgement

SZ and HZ are partially supported by an NSF IIS grant No. 2416897. HZ would like to thank the support from a Google Research Scholar Award. MY was supported by MEXT KAKENHI Grant Number 24K03004. The views and conclusions expressed in this paper are solely those of the authors and do not necessarily reflect the official policies or positions of the supporting companies and government agencies.

## Appendix

## Appendix A Approximation Methods of Earth Mover’s Distance

Most approximation methods do not rely on any tree structure. Sinkhorn [Cuturi, ] is a fully differentiable approximation method which leverages the entropic regularization setup of EMD and the convergence of doubly stochastic matrix. Sliced Wasserstein Distance (SWD) [Bonneel et al., ] projects data to 1 dimension through a random projection matrix and uses $\text{EMD}_{\text{1d}}$ for fast calculation, but the quality of approximation highly relies on the number of random projections. Tree Wasserstein Distance (TWD) [Yamada et al., , Le et al., , Takezawa et al., ] is a tree based approximation algorithm that can be written in the closed matrix form $\left\lVert\operatorname{diag}(\mathrm{w})\bm{B}(a-b)\right\rVert_{1}$, where $\mathrm{w}$ contains weights of ${\mathcal{T}}$ and ${\bm{B}}$ is a 0-1 internal nodes by leaf nodes matrix that represents if a node is the ancestor of a leaf node. To make sure TWD is a good approximation, an extra step of lasso regression is performed to learn $\mathrm{w}$ such that the tree metric approximates Euclidean distance between features with iterative gradient based method. In Yamada et al. [], it requires us to learn the tree distance between all pairs of nodes, which leads to quadratic computational complexity with respect to sample size as shown in Tab.1.

## Appendix B EMD-CPCC Generalizes ℓ 2 subscript ℓ 2 \ell_{2} roman_ℓ start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT -CPCC

We can show EMD-CPCC is a generalization of previous work of $\ell_{2}$-CPCC through the following proposition.
See 3.1

###### Proof.

Since $\mu,\nu$ in EMD definition are probability measures, the definition of EMD can be generalized to continuous distributions as well. Denote $Z$ the random variable from $\mu$, $Z^{\prime}$ from $\nu$, and $q$ the coupling between $\mu$ and $\nu$, then EMD can be also expressed in the following $p$-Wasserstein (${\mathcal{W}}_{p}$) form when $p=1$:

$$ $\displaystyle{\mathcal{W}}_{p}$ $\displaystyle=\left(\inf_{q}\int_{{\mathcal{Z}}\times{\mathcal{Z}}}\left\lVert z _{i}-z_{j}\right\rVert^{p}q(z_{i},z_{j})dz_{i}z_{j}\right)^{\frac{1}{p}}$ $\displaystyle=\inf_{q}(\mathbb{E}\lVert Z-Z^{\prime}\rVert^{p})^{\frac{1}{p}}$ (4) $$

With this definition, our statement can be proved in two steps. First, we want to show the lower bound of EMD is $\left\lVert\mu_{z}-\mu_{z^{\prime}}\right\rVert$ for any pairs of input random variables. By one step application of Jensen’s inequality, since every norm is convex, we have the following result from the expectation form of ${\mathcal{W}}_{1}$ in Eq.4:

$$ $\displaystyle{\mathcal{W}}_{1}=\inf\mathbb{E}\left\lVert Z-Z^{\prime}\right\rVert$ $\displaystyle\geq\inf\left\lVert\mathbb{E}(Z-Z^{\prime})\right\rVert$ $\displaystyle=\left\lVert\mathbb{E}Z-\mathbb{E}Z^{\prime}\right\rVert=\left \lVert\mu_{z}-\mu_{z^{\prime}}\right\rVert$ (5) $$

Second, we need to show EMD $\leq\left\lVert\mu_{z}-\mu_{z^{\prime}}\right\rVert$ for the given Gaussian constraints. Following the convexity of quadratic function, apply Jensen’s again, we have:

$$ $\displaystyle{\mathcal{W}}_{1}^{2}=(\inf\mathbb{E}\left\lVert Z-Z^{\prime} \right\rVert)^{2}\leq\inf\mathbb{E}(\left\lVert Z-Z^{\prime}\right\rVert)^{2}= {\mathcal{W}}_{2}^{2}$ (6) $$

Dowson and Landau [] showed the following analytical form of ${\mathcal{W}}_{2}^{2}$: $\lVert\mu_{z}-\mu_{z^{\prime}}\rVert^{2}+\text{tr}\left(\Sigma_{z}+\Sigma_{z^{
\prime}}-2(\Sigma_{z}\Sigma_{z^{\prime}})^{1/2}\right)$. The trace term becomes 0 0 if $\Sigma_{z}=\Sigma_{z^{\prime}}=\Sigma$.

∎

###### Proof.

## Appendix C Details for Fast Flow Tree CPCC Algorithm

In this section, we first present the full pseudocode to compute FastFT-CPCC in Alg.3. Compared to FlowTree [Backurs et al., ] which uses the helper function in Alg.2 [Kalantari and Kalantari, ], we replace the greedy_flow_matching algorithm with Alg.1 which aligns with algorithm that computes the 1 dimensional OT transportation plan.

Figure: Algorithm 3 Train with Fast FlowTree CPCC

Now we prove the correctness of Fast Flow Tree algorithm and its relation to 1 dimensional optimal transport distance.

See 3.2

###### Proof.

Alg.2 will not start to match leaves until the current node has children from two different distributions, or in our context, from two different fine labels. This is because of the while statement in line $8$: based on our construction of $\mathcal{T}$ in Fig.2, if the current node $v$ is either a leaf or a fine label node, then all of its children are from either $C_{a}(v)$ or $C_{b}(v)$, so we will jump to line $11$ to send all flows to $v$’s parent. Therefore, the matching starts at the lowest common ancestor of two label nodes.

Now we have $v$ as the lowest common ancestor of two fine class nodes (ex., $F$ or $H$ in Fig.2). Before entering line $9$, we always pass unmatched demands to the parent, and we can see the demands are always from one label until we reach the lowest common ancestor. From Fig.2, we also know at this point, we have collected all flows from both labels. Therefore, the sum of flow is $1$ for both fine label nodes when the matching starts.

Both Alg.2 and Alg.1 greedily remove one sample from either distribution. The greedy plan is always feasible and no unmatched flow is left after one scan of all $m+n$ features, because we inherently remove the the same amount of flow from both distributions every iteration in the while loop. Thus, on our ${\mathcal{T}}$, we can skip most steps in Alg.2: bottom up step in line $2$, looping through all vertices at height $h$, the if-else check from line $4$-$7$, and the last step of passing unmatched demands to parents are all unnecessary. Alg.2 terminates immediately once we run through all while loop steps, thus reducing to Alg.1.

Alg.1, or the flow matrix solution for 1-$d$ EMD problem, can be used when Monge Property or north-west-corner rules [Hoffman, ] hold. It can be shown by induction that this $O(m+n)$ greedy algorithm can be applied whenever the following constraint is satisfied: if $\{(z_{i},a_{i})\},\{(z_{j},b_{j})\}$ are indexed (1d feature, weight) sequences, for indices $i_{1}<i_{2}$ and $j_{1}<j_{2}$:

$$ ${\bm{D}}_{i_{1},j_{1}}+{\bm{D}}_{i_{2},j_{2}}\leq{\bm{D}}_{i_{1},j_{2}}+{\bm{D }}_{i_{2},j_{1}}$ (7) $$

If two sequences are ordered (in ascending order) by metric defining ${\bm{D}}$, we have $z_{i_{1}}\leq z_{i_{2}}$ and $z_{j_{1}}\leq z_{j_{2}}$.
There are in total six cases for any $i_{1}<i_{2},j_{1}<j_{2}$:

- 1.
If $z_{j_{1}}\leq z_{j_{2}}\leq z_{i_{1}}\leq z_{i_{2}}$, then Eq.7 $\Leftrightarrow\cancel{z_{i_{1}}}\cancel{-z_{j_{1}}}\cancel{+z_{i_{2}}}\cancel
{-z_{j_{2}}}\leq\cancel{z_{i_{1}}}\cancel{-z_{j_{2}}}\cancel{+z_{i_{2}}}
\cancel{-z_{j_{1}}}$ holds.
- 2.
If $z_{j_{1}}\leq z_{i_{1}}\leq z_{j_{2}}\leq z_{i_{2}}$, then Eq.7 $\Leftrightarrow z_{i_{1}}\cancel{-z_{j_{1}}}\cancel{+z_{i_{2}}}-z_{j_{2}}\leq z
_{j_{2}}-z_{i_{1}}\cancel{+z_{i_{2}}}\cancel{-z_{j_{1}}}$ holds because $z_{i_{1}}\leq z_{j_{2}}$.
- 3.
If $z_{j_{1}}\leq z_{i_{1}}\leq z_{i_{2}}\leq z_{j_{2}}$, then Eq.7 $\Leftrightarrow z_{i_{1}}\cancel{-z_{j_{1}}}\cancel{+z_{j_{2}}}-z_{i_{2}}\leq
\cancel{z_{j_{2}}}-z_{i_{1}}+z_{i_{2}}\cancel{-z_{j_{1}}}$ holds because $z_{i_{1}}\leq z_{i_{2}}$.
- 4.
If $z_{i_{1}}\leq z_{j_{1}}\leq z_{j_{2}}\leq z_{i_{2}}$, then Eq.7 $\Leftrightarrow z_{j_{1}}\cancel{-z_{i_{1}}}\cancel{+z_{j_{2}}}-z_{i_{2}}\leq
\cancel{z_{j_{2}}}\cancel{-z_{i_{1}}}+z_{i_{2}}-z_{j_{1}}$ holds because $z_{j_{1}}\leq z_{i_{2}}$.
- 5.
If $z_{i_{1}}\leq z_{j_{1}}\leq z_{i_{2}}\leq z_{j_{2}}$, then Eq.7 $\Leftrightarrow z_{j_{1}}\cancel{-z_{i_{1}}}\cancel{+z_{j_{2}}}-z_{i_{2}}\leq
\cancel{z_{j_{2}}}\cancel{-z_{i_{1}}}+z_{i_{2}}-z_{j_{1}}$ holds because $z_{j_{1}}\leq z_{i_{2}}$.
- 6.
If $z_{i_{1}}\leq z_{i_{2}}\leq z_{j_{1}}\leq z_{j_{2}}$, then Eq.7 $\Leftrightarrow\cancel{z_{j_{1}}}\cancel{-z_{i_{1}}}\cancel{+z_{j_{2}}}\cancel
{-z_{i_{2}}}\leq\cancel{z_{j_{2}}}\cancel{-z_{i_{1}}}\cancel{+z_{j_{1}}}
\cancel{-z_{i_{2}}}$ holds.

Therefore, if two sequences are ordered with respect to some metric, Eq.7 is satisfied.

Recall we need the extra sorting step in the beginning of 1d OT plan to apply the greedy flow matching algorithm. The distance function in Eq.7 is decided by the ground metric of EMD, which does not make any difference for 1 dimensional features. In Flow Tree, a tree is constructed with QuadTree such that the tree distance of any two features is close to the $\ell_{2}$ distance between them. In this case, in EMD computation, we can improve computational efficiency by Flow Tree. With extended label tree ${\mathcal{T}}$, because all same class sample leaves belong to one single parent, we have a special case here where the tree distance of any two samples from $\mu,\nu$ are always the same, satisfying Eq.7 vacuously. This statement fails only if we assign different edge weights for each sample during the extension of label tree, and does not depend on the weights of original label tree. To avoid the limitation on the extended edge weights, we can instead assign different weights $a_{i}$ for each sample to adjust the importance of each dat points. Therefore, FastFT is still very flexible.

For the time complexity, Alg.1 requires at most $O(m+n)$ steps from line $4$. The most time-consuming aspect is computing the $d$-dimensional distance of $O(m+n)$ pairs in ${\bm{D}}$, since there will be at most $m+n$ nonzero values in ${\bm{D}}$ after applying the greedy matching algorithm, which makes the final time complexity as $O((m+n)d)$.
∎

Figure: (a) Step 1
Refer to caption: x10.png

###### Proof.

## Appendix D Dataset Hierarchy Construction

In this section, we discuss the details of dataset hierarchies and how we created the dataset splits. Throughout the main body of this paper, we consider a two-level hierarchy: coarse and fine. For each coarse node, except for BREEDS, we split the coarse node’s fine children into approximately 40% for the target set and 60% for the source set. Following this procedure, the splits are detailed in Tables 6, 7, 8, 9, 10, 11, and 12. The original class labels for CIFAR10 and CIFAR100 are from [https://www.cs.toronto.edu/~kriz/cifar.html](https://www.cs.toronto.edu/~kriz/cifar.html), for iNaturalist from [https://github.com/visipedia/inat_comp/tree/master/2021](https://github.com/visipedia/inat_comp/tree/master/2021), and for BREEDS from WordNet [Miller, ].

The CIFAR10 hierarchy was manually constructed by us, while the CIFAR100 hierarchy is provided in the original dataset. In INAT, fine and coarse classes correspond to the class and phylum labels from the original annotations. These higher-level labels are particularly suitable for evaluating multi-modal settings, where OT-CPCC methods may have a more significant advantage over $\ell_{2}$-CPCC.

In Zeng et al. [], fine classes in BREEDS are defined one level deeper than the original labels. However, due to the large number of classes and limited samples per class within a batch, the distinction between $\ell_{2}$-CPCC and OT-CPCC can become less apparent. To address this, we employ a three-level hierarchy, using the labels from Santurkar et al. [] as fine labels and setting their parent categories as mid and coarse labels. Consequently, due to BREEDS setting, the fine classes are the same for both source and target datasets, so there is no source-target split in Tables 9, 10, 11, and 12. Despite this hierarchical structure, as seen in Fig.6, BREEDS exhibits relatively low multi-modality. With this design, in Sec.4.3, for classification and retrieval experiments, one-shot fine-tuning is not required for BREEDS on the target split.

**Table 6: CIFAR10 hierarchy and data split.**
| Coarse | Source | Target |
| --- | --- | --- |
| animal | deer, dog, frog | bird, cat |
| transportation | ship, truck | airplane, automobile |

**Table 7: CIFAR100 hierarchy and data split.**
| Coarse | Source | Target |
| --- | --- | --- |
| aquatic mammals | otter, seal, whale | beaver, dolphin |
| fish | ray, shark, trout | aquarium fish, flat fish |
| flowers | roses, sunflowers, tulips | orchids, poppies |
| food containers | cans, cups, plates | bottles, bowls |
| fruit and vegetables | oranges, pears, sweet pepper | apples, mushrooms |
| household electrical devices | lamp, telephone, television | clock, computer keyboard |
| household furniture | couch, table, wardrobe | bed, chair |
| insects | butterfly, caterpillar, cockroach | bee, beetle |
| large carnivores | lion, tiger, wolf | bear, leopard |
| large man-made outdoor things | house, road, skyscraper | bridge, castle |
| large natural outdoor scenes | mountain, plain, sea | cloud, forest |
| large omnivores and herbivores | chimpanzee, elephant, kangaroo | camel, cattle |
| medium-sized mammals | possum, raccoon, skunk | fox, porcupine |
| non-insect invertebrates | snail, spider, worm | crab, lobster |
| people | girl, man, woman | baby, boy |
| reptiles | lizard, snake, turtle | crocodile, dinosaur |
| small mammals | rabbit, shrew, squirrel | hamster, mouse |
| trees | palm, pine, willow | maple, oak |
| vehicles 1 | motorcycle, pickup truck, train | bicycle, bus |
| vehicles 2 | streetcar, tank, tractor | lawn-mower, rocket |

**Table 8: INAT hierarchy and data split.**
| Coarse | Source | Target |
| --- | --- | --- |
| Arthropoda | Diplopoda, Hexanauplia, Insecta, Malacostraca, Merostomata | Arachnida, Chilopoda |
| Chordata | Ascidiacea, Aves, Elasmobranchii, Mammalia, Reptilia | Actinopterygii, Amphibia |
| Cnidaria | Hydrozoa, Scyphozoa | Anthozoa |
| Echinodermata | Echinoidea, Holothuroidea, Ophiuroidea | Asteroidea |
| Mollusca | Cephalopoda, Gastropoda, Polyplacophora | Bivalvia |
| Ascomycota | Lecanoromycetes, Leotiomycetes, Pezizomycetes, Sordariomycetes | Arthoniomycetes, Dothideomycetes |
| Basidiomycota | Dacrymycetes, Pucciniomycetes, Tremellomycetes | Agaricomycetes |
| Bryophyta | Polytrichopsida, Sphagnopsida | Bryopsida |
| Tracheophyta | Liliopsida, Lycopodiopsida, Magnoliopsida, Pinopsida, Polypodiopsida | Cycadopsida, Gnetopsida |

**Table 9: LIVING17 dataset hierarchy.**
| Coarse | Fine |
| --- | --- |
| bird | grouse, parrot |
| amphibian | salamander |
| reptile, reptilian | turtle, lizard, snake, serpent, ophidian |
| arthropod | spider, crab, beetle, butterfly |
| mammal, mammalian | dog, domestic dog, Canis familiaris, wolf, fox, domestic cat, house cat, Felis domesticus, Felis catus, bear, ape, monkey |

**Table 10: NONLIVING26 dataset hierarchy.**
| Coarse | Fine |
| --- | --- |
| structure, place | dwelling, home, domicile, abode, habitation, dwelling house, fence, fencing, mercantile establishment, retail store, sales outlet, outlet, outbuilding, roof |
| paraphernalia | ball, bottle, digital computer, keyboard instrument, percussion instrument, percussive instrument, pot, stringed instrument, timepiece, timekeeper, horologe, wind instrument, wind |
| food, nutrient | squash |
| apparel, toiletries | bag, body armor, body armour, suit of armor, suit of armour, coat of mail, cataphract, coat, hat, chapeau, lid, skirt |
| conveyance, transport | boat, bus, autobus, coach, charabanc, double-decker, jitney, motorbus, motorcoach, omnibus, passenger vehicle, car, auto, automobile, machine, motorcar, ship, truck, motortruck |
| furnishing | chair |

**Table 11: ENTITY13 dataset hierarchy.**
| Coarse | Fine |
| --- | --- |
| living thing, animate thing | bird, reptile, reptilian, arthropod, mammal, mammalian |
| non-living thing | garment, accessory, accoutrement, accouterment, craft, equipment, furniture, piece of furniture, article of furniture, instrument, man-made structure, construction, wheeled vehicle, produce, green goods, green groceries, garden truck |

**Table 12: ENTITY30 dataset hierarchy.**
| Coarse | Fine |
| --- | --- |
| living thing, animate thing | serpentes, passerine, passeriform bird, saurian, arachnid, arachnoid, aquatic bird, crustacean, carnivore, insect, ungulate, hoofed mammal, primate, bony fish |
| non-living thing | barrier, building, edifice, electronic equipment, footwear, legwear, garment, headdress, headgear, home appliance, household appliance, kitchen utensil, measuring instrument, measuring system, measuring device, motor vehicle, automotive vehicle, musical instrument, instrument, neckwear, sports equipment, tableware, tool, vessel, watercraft, dish, vegetable, veggie, veg, fruit |

## Appendix E Training Hyperparameters

We conducted all experiments with NVIDIA RTX A6000. The architectural and training setups for each dataset are as follows:

### E.1 CIFAR10, CIFAR100

We conduct training from scratch using ResNet18, following the ResNet architecture design where the first convolution layer employs a 3x3 kernel size, while the max pooling layer is omitted. The images are preprocessed by normalizing them according to CIFAR10/CIFAR100’s mean and standard deviation. We augment the training data by first zero-padding the images by 4 pixels, then cropping them back to their original dimensions, applying a 50% chance of horizontal flipping, and randomly rotating the images between -15 and 15 degrees.

For CIFAR10, the pretraining process spans 100 epochs with a batch size of 64. The learning rate begins at 0.01 and is decreased by a factor of 10 every 60 steps, with momentum set at 0.9 and weight decay set to $5\times 10^{-4}$ in the SGD optimizer. For CIFAR100, the pretraining process spans 200 epochs with a batch size of 128. The learning rate begins at 0.1 and is decreased by a factor of 5 every 60 steps, with momentum set at 0.9 and weight decay set to $5\times 10^{-4}$ in the SGD optimizer. Additionally, for both datasets, in the fine-level one-shot transfer learning, we train for 100 epochs starting at a learning rate of 0.1, which is reduced by a factor of 10 every 60 steps. CPCC is applied to coarse, fine level and $\lambda$ in Eq.1 is set to $1$. The ratio between the distance between two fine classes sharing the same coarse parent and those sharing different parents is 1:2.

### E.2 BREEDS

We train a ResNet18 model from scratch originally designed for ImageNet, where the first convolution layer uses a 7x7 kernel. Adhering to the BREEDS training setup, we use a batch size of 128, an initial learning rate of 0.1, and SGD optimization with a weight decay of $10^{-4}$. For the ENTITY13 and ENTITY30 datasets, the model is trained for 300 epochs, reducing the learning rate by a factor of 10 every 100 steps, while for LIVING17 and NONLIVING26, we extend training to 450 epochs, with the learning rate divided by 10 every 150 steps. Data augmentation for the training set includes a RandomResizedCrop to 224x224, horizontal flipping with a 50% probability, and color jittering that adjusts brightness, contrast, and saturation within a range of 0.9 to 1.1. The test set, on the other hand, is resized to 256x256 and then center-cropped to 224x224.

Since in our experiments with BREEDS, the source and target split share the same set of fine labels, to calculate fine target accuracy, it is unnecessary to use one-shot transfer learning. We directly use the source set pretrained models for target fine accuracy evaluation.

To apply CPCC methods, we backtrack two levels up to find the mid and coarse labels from the original class labels in BREEDS. CPCC is applied to mid and coarse level and $\lambda$ in Eq.1 is set to $1$. The ratio between the distance between two fine classes sharing the same coarse parent and those sharing different parents is 1:2.

### E.3 INAT

We fine-tune an ImageNet-pretrained ResNet18 model where checkpoint can be downloaded in Pytorch (torchvision.models.resnet18(weights=’IMAGENET1K_V1’), [https://pytorch.org/vision/main/models/generated/torchvision.models.resnet18.html](https://pytorch.org/vision/main/models/generated/torchvision.models.resnet18.html)). The images are preprocessed by normalizing them with mean=[0.466, 0.471, 0.380] and standard deviation=[0.195, 0.194, 0.192]. We augment the training data by cropping a random portion (scale=(0.08, 1.0), ratio=(0.75, 1.33)) of image and resize it to 224x224, and apply a 50% chance of horizontal flipping. We use a batch size of 1024. The learning rate begins at 0.1 and is decreased by a factor of 5 every 60 steps, with momentum set at 0.9 and weight decay set to $5\times 10^{-4}$ in the SGD optimizer. In the downstream one-shot fine-tuning task, we train for 60 epochs starting at a learning rate of 0.05, which is reduced by a factor of 10 every 15 steps. CPCC is applied to coarse, fine level and $\lambda$ in Eq.1 is set to $0.1$. The ratio between the distance between two fine classes sharing the same coarse parent and those sharing different parents is 1:3.

Figure: Figure 9: Visualization of the evaluation of hierarchical classification.
Refer to caption: x22.png

## Appendix F Hierarchical Baseline Objectives

In this section, we introduce the details of hierarchical objectives we used in Fig.12.

#### Flat

is the naive cross entropy loss only on the fine level, so the expression is:

$$ $\mathcal{L}_{\textnormal{Flat}}(x,y_{\textnormal{fine}})=\sum\limits_{(x,y)\in \mathcal{D}}\ell_{\textnormal{Flat}}(y_{\textnormal{fine}},h(x)).$ (8) $$

We adopt all hierarchical baselines from the comparative study in Zeng et al. [], including:

Multi-Task Learning (MTL) [Caruana, ], where the backbone model is trained with two linear heads—one for coarse classes and one for fine classes—each with a cross-entropy loss. The total loss, $\mathcal{L}_{\textnormal{MTL}}$, is the sum of these two cross-entropy terms, with a scaling coefficient of 1 for both. Curriculum Learning (Curr) [Bengio et al., ] involves two training phases: first, training a coarse head for 50 epochs, followed by training the full network with a fine head for 200 epochs, starting from the coarse-trained features, assuming that coarse classification is an easier task. Sum Loss [Hu et al., ] uses only one fine classifier head, where coarse predictions are obtained by summing the fine class probabilities for the same coarse parent and applying a second cross-entropy loss. Hierarchical Cross Entropy (HXE) [Bertinetto et al., ] is the sum of weighted cross-entropy losses ($\alpha=0.4$), where the hierarchical structure is incorporated into the log-probabilities. The same paper also introduces Soft Label (Soft), which uses a single softmax loss based on the height of the least common ancestor between label pairs, with a hyperparameter $\beta=10$. Quadruplet Loss (Quad) [Zhang et al., ] extends triplet loss by defining two types of positive labels based on the hierarchy. For detailed mathematical expressions of these objectives, refer to App.C in Zeng et al. [].

We also include several recently published hierarchical losses:

#### Soft Hierarchical Consistency

(SHC) [Garg et al., ] is the combination of MTL and Sum Loss. For a two level hierarchy, they jointly train a two-headed network. Define the network as $h:\mathcal{X}\rightarrow\Delta_{k_{1}}\times\Delta_{k_{2}}$ with a shared feature encoder $f:\mathcal{X}\rightarrow\mathcal{Z}$ and two classifiers $g_{\textnormal{coarse}}:\mathcal{Z}\rightarrow\Delta_{k_{1}},g_{\textnormal{
fine}}:\mathcal{Z}\rightarrow\Delta_{k_{2}}$, with $k_{1}$ as the number of coarse classes and $k_{2}$ the number of fine classes. The ground truth coarse probability $p_{\text{coarse}}=g_{\textnormal{coarse}}(f(\cdot))$ is the output of coarse head after softmax normalization, and $\hat{p}_{\text{coarse}}$ is derived from the output of fine head. Let $\mathbf{W}$ be a $k_{1}$ by $k_{2}$ matrix representing the relationships in the label tree: if a fine class $i$ belongs to a coarse class $j$, then $\mathbf{W}_{ji}$ is $1$, else the entry is set to 0 0. Then $\hat{p}_{\text{coarse}}=\mathbf{W}g_{\text{fine}}(f(\cdot))$, so the probability of belonging to a given coarse class is the sum of the probabilities of its fine class descendants in the tree. Then, the SHC loss is defined as the Jensen-Shannon Divergence between these two probability distributions:

$$ $\mathcal{L}_{\text{SHC}}(x,y_{\textnormal{coarse}},y_{\textnormal{fine}})= \textnormal{JSD}(p_{\text{coarse}}(x,y_{\textnormal{coarse}})||\hat{p}_{\text{ coarse}}(x,y_{\textnormal{fine}})).$ (9) $$

In Garg et al. [], $\mathcal{L}_{\text{SHC}}$ is used in the combination with Flat cross entropy loss. Therefore, we treat it as a counterpart of CPCC losses and use $\mathcal{L}_{\textnormal{Flat}}+\lambda\cdot\mathcal{L}_{\textnormal{SHC}}$ as the training loss, with $\lambda=1$.

#### Geometric Consistency

(GC) [Garg et al., ] is also computed under the MTL two-head pipeline. The coarse classifier can be decomposed to a linear layer and softmax function, i.e., let $\mathbf{W}^{c}\in\mathbb{R}^{d\times k_{1}}$, $g_{\text{coarse}}(\cdot)=\textnormal{softmax}(\mathbf{W}^{c}(\cdot))$, and we define $\mathbf{W}^{f}$ similarly for the linear fine layer.

$$ $\mathcal{L}_{\textnormal{GC}}(x,y_{\textnormal{coarse}},y_{\textnormal{fine}}) =\sum\limits_{i=1}^{k_{1}}1-\cos{(\mathbf{W}^{c}_{i},\hat{\mathbf{W}}^{c}_{i})}.$ (10) $$

As in SHC, GC aims to force the similarity between the predicted coarse weight and the true coarse weight. Here, $\mathbf{W}^{c}_{i}$ is part of the linear classifier corresponds to each coarse class, and the predicted coarse weight $\hat{\mathbf{W}}^{c}_{i}=\sum_{j\in\textnormal{children}(i)}\mathbf{W}^{f}_{j}
/\left\lVert\sum_{j\in\textnormal{children}(i)}\mathbf{W}^{f}_{j}\right\rVert$ is the normalized sum of a part of fine classifiers based on the hierarchical relationship. We combine $\mathcal{L}_{\text{GC}}$ with Flat cross entropy during training and set $\lambda$ in Eq.1 as 1.

#### RankLoss

(Rank) [Nolasco and Stowell, ] proposes a ranking loss to make Euclidean distance of feature $\rho$ close to the target $d$-dim sequence determined by the rank of tree metric:

$$ $\mathcal{L}_{\text{rank}}(x,y_{\text{coarse}},y_{\text{fine}})=\frac{1}{p}\sum _{p}(1-I_{p})(\rho-\text{TargetRank}(\rho)).$ (11) $$

In this equation above, $I_{p}$ is an indicator variable which shows whether the rank of pairwise $\ell_{2}$ feature distance is the same as the rank of pairwise tree metric distance, and TargetRank is a function reflecting hierarchical similarity. Since this loss does not focus on discriminative tasks like classification, we combine it with Flat cross entropy, setting $\lambda=1$ in Eq.1 for downstream tasks evaluation.

## Appendix G Additional Experiments

### G.1 Effect of Tree Depth

Figure: Figure 10: CIFAR100 class embedding pairwise distance matrix for different types of label tree with varying tree depth.
Refer to caption: extracted/5893798/images/CIFAR100_distM_comparison-depth.png

Although in the main body we focus on a depth-3 label hierarchy consisting of a root node, a coarse level, and a fine level, we also test whether our method can be applied to more complex trees. In CIFAR100, beyond the standard fine-coarse levels, we introduce a mid level by grouping 5 fine classes from the same coarse class into two smaller groups (2 + 3). Additionally, we create a coarser level by grouping two alphabetically neighboring coarse classes into a single coarser label. For example, for 10 fine labels, the hierarchical grouping is highlighted by the following color scheme: {[(orchids, poppies) (roses, sunflowers, tulips)] [(bottles, bowls) (cans, cups, plates)]}.

In Fig.10, we aligned the axes to sort the fine labels based on the hierarchical structure for embeddings learned using trees of varying depth. We then computed the fine-to-fine class distances using either the EMD or FastFT metric. On the left of Fig.10, we used the standard fine-coarse CIFAR100 label tree, where a clear pattern of twenty 5x5 coarse blocks appears on the diagonal for both optimal transport methods. In the middle of Fig.10, we applied CPCC to the fine-mid-coarse hierarchy, and within each 5x5 coarse block, we observed the expected fine-grained 2x2 + 3x3 patterns. Similarly, on the right of Fig.10, when we included the coarser grouping, we could observe the merging of neighboring 5x5 coarse blocks.

In summary, EMD-CPCC is a flexible structured representation learning method that adapts well to complex hierarchies, while FastFT-CPCC serves as a good approximation to EMD-CPCC. This observation is consistent with the results in Tab.3, where optimal transport methods achieve higher Test CPCC values.

### G.2 Effect of Tree Weights

Figure: Figure 11: CIFAR100 class embedding pairwise distance matrix for different types of label tree with varying tree weights.
Refer to caption: extracted/5893798/images/CIFAR100_distM_comparison-weight.png

In this section, we demonstrate that all CPCC methods, including OT-CPCC, are influenced by input tree weights. In Fig.11, we begin with the original unweighted CIFAR100 tree (where all edges have a default weight of 1) and modify the edge weights of a subtree representing the coarse label aquatic_mammals, which includes five fine labels: beaver, dolphin, otter, seal, and whale. We visualize the learned representation through a distance matrix, as described in Sec.G.1. From left to right in Fig.11, the edge weights of this subtree are set to 1, 2, and 4, respectively. Consequently, we observe a color shift in the top-left coarse block of the matrix, with the block gradually darkening as the subtree weights increase. This demonstrates the flexibility of OT-CPCC methods, which can successfully embed weighted trees, and highlights why CPCC is a sensitive and effective metric for evaluating the interpretability and quality of the hierarchical information embedded in the learned representation, as discussed in Sec.4.4.

### G.3 Effect of Flow Weight

In Eq.3, the flow weight vectors $a$ and $b$ are set to uniform by default in all other experiments. However, in practice, these vectors can be customized by the user if prior information about the data distribution is available, which may significantly influence the learned embeddings. On CIFAR100, we tested three different flow weight settings:

- •
Uniform (unif): default setting where we assume each data point is uniformally sampled from the class distribution, so $a=(1/m,\dots,1/m)$ and $b=(1/n,\dots,1/n)$.
- •
Distance (dist): is an adaptive weighting method where during training, we compute the Euclidean centroid of each class and compute the distance of each sample to the corresponding class centroid. Then we can apply softmax function to normalize the distances into probabilities. Mathematically, if two class representations are $Z,Z^{\prime}$, then for each entry of flow weight vector, $a_{i}=\frac{\exp(Z_{i}-\mu_{z})}{\sum_{j=1}^{m}\exp(Z_{j}-\mu_{z})}$, $b_{i}=\frac{\exp(Z_{i}^{\prime}-\mu_{z^{\prime}})}{\sum_{j=1}^{n}\exp(Z_{j}^{
\prime}-\mu_{z^{\prime}})}$.
- •
Inverse-Distance (inv): is similar to dist, yet the only difference is we set the weight to be inversely scaled by the distance to the centroid, so the expressions of the flow weights are $a_{i}=\frac{\exp(\mu_{z}-Z_{i})}{\sum_{j=1}^{m}\exp(\mu_{z}-Z_{j})}$, $b_{i}=\frac{\exp(\mu_{z^{\prime}}-Z_{i}^{\prime})}{\sum_{j=1}^{n}\exp(\mu_{z^{
\prime}}-Z_{j}^{\prime})}$.

From Tab.13, we observe two key findings: first, tuning the flow weights can lead to improvements in out-of-distribution generalization metrics for the target set, but this comes at the expense of in-distribution source metrics. Second, the most significant differences are seen in the CPCC values, where it is nearly impossible to recover hierarchical information on the test set when using the Distance flow weight for both EMD and FastFT. The low interpretability scores suggest that not all flow weight settings are suitable for hierarchy learning.

**Table 13: Comparison of FastFT and EMD methods with different optimal transport flow weighting schemes across fine and coarse metrics, with cell colors representing the magnitude of values.**
| Objective | Weight | CPCC | Fine Metrics | Coarse Metrics |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| sAcc | sMAP | tAcc | sAcc | sMAP | tAcc | tMAP |  |  |  |
| FastFT | unif | 93.85 | 81.32 | 78.88 | 23.15 | 86.96 | 91.71 | 44.29 | 42.20 |
| dist | 9.30 | 78.90 | 67.33 | 25.28 | 85.18 | 91.47 | 45.00 | 43.78 |  |
| inv | 66.52 | 79.73 | 68.31 | 25.45 | 85.81 | 92.03 | 43.70 | 43.70 |  |
| EMD | unif | 89.60 | 80.99 | 78.51 | 23.53 | 86.93 | 91.82 | 44.03 | 42.32 |
| dist | 19.87 | 79.23 | 73.79 | 26.05 | 85.91 | 91.30 | 43.55 | 42.53 |  |
| inv | 98.13 | 79.51 | 75.84 | 26.80 | 86.05 | 91.22 | 44.15 | 42.13 |  |

**Table 14: Detailed generalization performance on fine level. Each value is the mean (std) of three seeds of experiments. Flat directly optimizes on the fine level classes, so we treat it as the upper bound of fine metrics and colored in grey. The bold values have the largest mean for a column metric within a dataset, ignoring the upper bound.**
| Dataset | Objective | sAcc | tAcc | sMAP | Dataset | sAcc | tAcc | sMAP |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| CIFAR10 | Flat | 96.92 (0.14) | 54.23 (6.92) | 99.43 (0.02) | INAT | 88.70 (0.01) | 26.77 (0.04) | 63.97 (0.04) |
| $\ell_{2}$ | 96.96 (0.16) | 55.71 (7.38) | 99.22 (0.06) | 88.62 (0.01) | 26.78 (0.04) | 56.10 (0.05) |  |  |
| FastFT | 96.90 (0.13) | 55.99 (8.19) | 99.24 (0.02) | 88.49 (0.01) | 27.10 (0.04) | 56.21 (0.05) |  |  |
| EMD | 97.05 (0.18) | 56.12 (7.63) | 99.24 (0.02) | 88.68 (0.00) | 26.78 (0.04) | 56.83 (0.00) |  |  |
| Sinkhorn | 96.95 (0.08) | 54.89 (8.59) | 99.27 (0.03) | 88.08 (0.00) | 26.77 (0.04) | 51.56 (0.04) |  |  |
| SWD | 96.96 (0.05) | 59.21 (4.42) | 99.45 (0.03) | 88.46 (0.00) | 26.78 (0.04) | 54.19 (0.03) |  |  |
| CIFAR100 | Flat | 80.53 (0.31) | 28.44 (1.89) | 86.47 (0.21) | BREEDS | 82.08 (0.26) | 45.19 (0.38) | 73.74 (0.46) |
| $\ell_{2}$ | 80.98 (0.17) | 23.76 (1.01) | 77.52 (0.12) | 82.66 (0.41) | 45.95 (0.40) | 62.86 (0.38) |  |  |
| FastFT | 81.09 (0.18) | 24.71 (0.90) | 78.72 (0.19) | 82.58 (0.34) | 45.75 (0.45) | 63.95 (0.36) |  |  |
| EMD | 81.32 (0.19) | 23.15 (0.40) | 78.88 (0.18) | - | - | - |  |  |
| Sinkhorn | 80.99 (0.12) | 23.53 (0.48) | 78.51 (0.20) | 82.99 (0.31) | 46.87 (0.40) | 61.54 (0.55) |  |  |
| SWD | 80.50 (0.33) | 26.18 (1.04) | 86.82 (0.22) | 82.90 (0.30) | 46.33 (0.16) | 76.24 (0.39) |  |  |

**Table 15: Detailed generalization performance on coarse level. Each value is the mean (std) of three seeds of experiments. $\ell_{2}$ directly optimizes on the coarse level classes, so we treat it as the upper bound of coarse metrics and colored in grey.**
| Dataset | Objective | sAcc | tAcc | sMAP | tMAP | Dataset | sAcc | tAcc | sMAP | tMAP |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| CIFAR10 | $\ell_{2}$ | 99.60 (0.01) | 87.51 (0.97) | 99.88 (0.03) | 92.84 (0.63) | INAT | 94.52 (0.00) | 40.86 (0.01) | 74.34 (0.05) | 37.35 (0.01) |
| Flat | 99.58 (0.07) | 87.30 (0.59) | 99.22 (0.05) | 89.66 (0.38) | 94.63 (0.00) | 38.81 (0.01) | 70.41 (0.04) | 34.00 (0.13) |  |  |
| FastFT | 99.61 (0.06) | 87.79 (0.41) | 99.91 (0.03) | 93.04 (0.57) | 94.66 (0.00) | 39.43 (0.02) | 72.90 (0.05) | 35.63 (0.01) |  |  |
| EMD | 99.61 (0.03) | 87.45 (0.60) | 99.88 (0.03) | 93.07 (0.63) | 94.64 (0.00) | 41.01 (0.01) | 73.87 (0.01) | 35.58 (0.01) |  |  |
| Sinkhorn | 99.61 (0.05) | 87.41 (0.53) | 99.87 (0.03) | 92.60 (0.13) | 94.30 (0.00) | 36.94 (0.02) | 68.75 (0.02) | 34.92 (0.00) |  |  |
| SWD | 99.56 (0.04) | 87.63 (0.62) | 99.36 (0.09) | 90.12 (0.43) | 94.52 (0.00) | 39.38 (0.01) | 75.13 (0.03) | 38.11 (0.01) |  |  |
| CIFAR100 | $\ell_{2}$ | 87.00 (0.13) | 44.55 (0.39) | 92.13 (0.23) | 43.03 (0.55) | BREEDS | 94.24 (0.30) | 83.13 (0.23) | 95.00 (0.25) | 81.73 (0.32) |
| Flat | 86.17 (0.26) | 43.09 (0.55) | 82.07 (0.34) | 42.09 (0.26) | 93.85 (0.20) | 82.43 (0.27) | 74.47 (0.52) | 66.50 (0.39) |  |  |
| FastFT | 86.96 (0.13) | 44.29 (0.33) | 91.71 (0.11) | 42.20 (0.44) | 94.25 (0.27) | 82.89 (0.31) | 92.84 (0.27) | 79.69 (0.34) |  |  |
| EMD | 86.93 (0.28) | 44.03 (0.31) | 91.82 (0.14) | 42.32 (0.31) | - | - | - | - |  |  |
| Sinkhorn | 86.93 (0.10) | 43.99 (0.37) | 91.61 (0.31) | 42.45 (0.09) | 94.42 (0.21) | 83.51 (0.30) | 91.84 (0.42) | 79.33 (0.38) |  |  |
| SWD | 86.42 (0.26) | 43.35 (0.71) | 87.24 (0.29) | 42.39 (0.30) | 94.41 (0.25) | 82.96 (0.17) | 83.58 (0.60) | 72.38 (0.43) |  |  |

**Table 16: Fine level generalization performance for each BREEDS setting.**
| Dataset | Objective | sAcc | tAcc | sMAP |
| --- | --- | --- | --- | --- |
| L17 | Flat | 85.15 (0.30) | 46.31 (0.31) | 85.94 (0.18) |
| $\ell_{2}$ | 85.74 (0.69) | 47.50 (0.73) | 74.07 (0.61) |  |
| FastFT | 85.21 (0.38) | 46.47 (0.75) | 74.46 (0.16) |  |
| Sinkhorn | 85.80 (0.49) | 48.37 (0.69) | 70.50 (0.53) |  |
| SWD | 85.88 (0.26) | 47.50 (0.09) | 86.65 (0.51) |  |
| N26 | Flat | 79.76 (0.27) | 36.06 (0.37) | 70.75 (1.14) |
| $\ell_{2}$ | 80.58 (0.54) | 37.08 (0.26) | 55.51 (0.50) |  |
| FastFT | 80.78 (0.52) | 36.69 (0.19) | 58.87 (0.42) |  |
| Sinkhorn | 81.74 (0.26) | 38.16 (0.14) | 57.54 (0.92) |  |
| SWD | 81.79 (0.25) | 37.71 (0.09) | 76.03 (0.05) |  |
| E13 | Flat | 84.37 (0.24) | 56.04 (0.27) | 74.04 (0.34) |
| $\ell_{2}$ | 84.77 (0.18) | 55.69 (0.37) | 66.92 (0.16) |  |
| FastFT | 85.04 (0.15) | 56.66 (0.44) | 67.96 (0.64) |  |
| Sinkhorn | 85.00 (0.15) | 57.14 (0.40) | 64.47 (0.20) |  |
| SWD | 84.49 (0.18) | 56.82 (0.12) | 75.39 (0.54) |  |
| E30 | Flat | 79.05 (0.23) | 42.35 (0.57) | 64.23 (0.17) |
| $\ell_{2}$ | 79.54 (0.21) | 43.51 (0.23) | 54.92 (0.23) |  |
| FastFT | 79.29 (0.32) | 43.16 (0.42) | 54.49 (0.23) |  |
| Sinkhorn | 79.40 (0.33) | 43.82 (0.38) | 53.66 (0.53) |  |
| SWD | 79.42 (0.52) | 43.28 (0.34) | 66.87 (0.45) |  |

**Table 17: Coarse level generalization performance for each BREEDS setting.**
| Dataset | Objective | sAcc | tAcc | sMAP | tMAP |
| --- | --- | --- | --- | --- | --- |
| L17 | $\ell_{2}$ | 93.13 (0.62) | 79.03 (0.29) | 93.93 (0.54) | 69.13 (0.42) |
| Flat | 92.68 (0.39) | 78.35 (0.37) | 76.66 (0.26) | 56.13 (0.36) |  |
| FastFT | 93.33 (0.26) | 78.37 (0.47) | 92.11 (0.23) | 65.44 (0.24) |  |
| Sinkhorn | 93.11 (0.20) | 79.50 (0.43) | 89.14 (0.36) | 63.88 (0.32) |  |
| SWD | 92.96 (0.41) | 78.66 (0.26) | 83.92 (0.38) | 59.66 (0.21) |  |
| N26 | $\ell_{2}$ | 88.94 (0.42) | 70.24 (0.50) | 89.59 (0.26) | 68.40 (0.52) |
| Flat | 88.44 (0.12) | 68.98 (0.41) | 61.20 (1.01) | 50.86 (0.16) |  |
| FastFT | 88.83 (0.75) | 69.43 (0.42) | 83.89 (0.40) | 65.22 (0.75) |  |
| Sinkhorn | 89.94 (0.35) | 70.70 (0.41) | 83.14 (0.78) | 66.07 (0.85) |  |
| SWD | 89.88 (0.34) | 69.93 (0.04) | 74.06 (0.57) | 58.50 (0.48) |  |
| E13 | $\ell_{2}$ | 97.80 (0.04) | 91.51 (0.01) | 98.25 (0.04) | 94.15 (0.09) |
| Flat | 97.53 (0.12) | 91.30 (0.16) | 83.89 (0.13) | 82.37 (0.31) |  |
| FastFT | 97.80 (0.01) | 91.78 (0.13) | 97.69 (0.13) | 93.21 (0.10) |  |
| Sinkhorn | 97.63 (0.02) | 91.97 (0.14) | 97.53 (0.23) | 92.77 (0.18) |  |
| SWD | 97.67 (0.13) | 91.69 (0.06) | 91.43 (0.72) | 87.71 (0.16) |  |
| E30 | $\ell_{2}$ | 97.07 (0.10) | 91.74 (0.10) | 98.24 (0.14) | 95.28 (0.24) |
| Flat | 96.75 (0.18) | 91.10 (0.15) | 76.12 (0.66) | 76.65 (0.74) |  |
| FastFT | 97.02 (0.07) | 91.98 (0.23) | 97.65 (0.32) | 94.90 (0.26) |  |
| Sinkhorn | 96.99 (0.25) | 91.87 (0.22) | 97.54 (0.31) | 94.59 (0.18) |  |
| SWD | 97.12 (0.13) | 91.55 (0.31) | 84.89 (0.74) | 83.66 (0.88) |  |

### G.4 Comparison to Other Hierarchical Baselines

In Fig.12, we compared OT-CPCC methods with other hierarchical objectives described in Sec.F, and we report the Fine/Coarse Metric Avg., which averages all source/target classification and retrieval metrics. Based on CPCC, fine, and coarse metrics, we observe that OT-CPCC methods behave similarly to $\ell_{2}$-CPCC, exhibiting higher CPCC values and stronger performance on coarse metrics, but lower fine metrics (though OT methods still outperform $\ell_{2}$ on fine metrics, as shown in Tab.5). In contrast, all other hierarchical baseline objectives perform more like flat cross-entropy methods, with the opposite trend—lower CPCC and coarse metrics but relatively higher fine metrics. Additionally, we excluded Rank loss from the correlation plot on the right side of Fig.12, as Tab.18 suggests that Rank loss is not suitable for any of the downstream tasks used in our evaluation, making it an outlier compared to the other objectives.

Figure: Figure 12: Left: Comparison of CPCC for hierarchical baselines on CIFAR100. Right: CPCC against coarse and fine metrics average. Rank is an outlier, so its value is removed from the visualization.
Refer to caption: x23.png

**Table 18: Comparison of fine and coarse metrics for Flat vs. Rank Loss.**
| Objective | Fine Metrics | Coarse Metrics |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| sAcc | tAcc | sMAP | sAcc | tAcc | sMAP | tMAP |  |
| Flat | 80.53 | 28.44 | 86.47 | 86.17 | 43.09 | 82.07 | 42.09 |
| Rank | 77.89 | 20.48 | 74.39 | 83.41 | 39.25 | 46.60 | 24.31 |

## Appendix H Implementation Details for Synthetic Datasets Experiments

#### Computational Efficiency

We measured the computation time for different OT methods across datasets of varying sizes. This involved generating two random $128$-dimensional datasets of identical shape and recording the wall-clock time for a single distance computation. This process was repeated $10$ times, and the average time was reported. For illustrative purposes, we fixed the number of projections at $10$ for SWD, and the regularization parameter at $10$ for Sinkhorn, the same as in our experiment setting in Tab.4,5.

#### Approximation Error

We evaluated the error of different OT approximation methods, along with the $\ell_{2}$ distance of the mean of two datasets, in approximating EMD. Two different distribution scenarios were considered:

- •
Gaussian Distribution: Here, both distributions were created as Gaussian distributions with different means (1 and 4, respectively) and the same standard deviation of 1.
- •
Non-Gaussian Distribution: For this scenario, we generated the dataset from Gaussian mixtures. One dataset combined Gaussian distributions with means of 0 and 5, while the other merged Gaussian distributions with means of 2 and 3, all of them with a standard deviation of 1.

Other experiment settings were consistent with those detailed in the Computational Efficiency paragraph above.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: See  the Weakness section.
soundness: 3
strengths: - The paper employs optimal transport (OT) distance to measure the distributions of labels rather than relying solely on the distance between class means. By capturing distributional differences, this approach aims to provide a more nuanced measure of class similarity.

- To address the high computational complexity typically associated with OT distance, the authors implement a tree-based approximation method. This approach enables more efficient computation while maintaining the benefits of OT-based similarity measures.

- Experimental results demonstrate that accounting for the distributional nature of labels enhances performance in tasks such as classification and class representation learning. This distribution-based approach proves beneficial in capturing class relationships more effectively than mean-based methods.
summary: The paper aims to improve classification by incorporating similarities among classes. These semantic relationships between labels are introduced into the original problem as a regularizer. The paper assumes a hierarchical structure in the labels. Building on the work of Zeng et al. (2022), the authors replace Euclidean distances between class means with optimal transport (OT) distances between label distributions. To reduce the computational burden of OT distance, they use a modified version of a tree-based method called FlowTree, which relies on a predefined tree structure. The paper also introduces a differentiable method for training with their tree-based OT distance. Experimental results for the classification problem using the OT-based regularization method show improvements compared to the method in Zeng et al. (2022). However, the proposed FastFT regularization falls behind the performance of some other OT-based methods.
weaknesses: - Theorem 3.2 appears to be a key contribution, as it establishes that Algorithms 1 and 2 can approximate the Earth Mover's Distance (EMD). However, the proof is vague and leaves several critical points open to interpretation. I found it challenging to follow the complete argument leading to the stated approximation result. Including a detailed proof sketch could greatly enhance clarity and strengthen the manuscript.

- Similarly, while Lemma 3.4 provides the derivative of the EMD, the result itself is relatively straightforward. More importantly, it is unclear whether the same derivative result directly applies to FastTree. A formal proof confirming that this result holds for FastTree would significantly enhance the paper's contribution and rigor.

- Another issue arises with the reliance on an assumed ground-truth tree structure. In most cases, this ideal tree is unknown or impractical to obtain, even if the underlying assumptions hold. As a result, multiple approximated trees may be used instead. The paper does not adequately address how gradient computations would handle these varied approximations, which is an important practical consideration.

- Finally, a literature review on learning label representations, accompanied by illustrative examples, would provide valuable context and motivation for the study. This would help readers appreciate the importance of this approach within the broader scope of label representation learning.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. While different approximation methods can be used, it is not clear about the computation cost, compared to the baseline methods. 
 
2. In Table 3, the preciseness of hierachical information listed show that in most cases, the presented FastFT based method cannot beat the $\ell_2$ norm based method. How to interprete the results in Table 3? 

3. Using AIC to estimate the number of modality of the classes sounds interesting.  Is there any more intuitive evidence to support the existence of multi-mode in feature distribution? Or,  whether the identified multi-mode classes could be visualized by e.g., t-SNE or UMAP?
soundness: 3
strengths: + It extends the $\ell_2$ norm based Cophenetic Correlation Coefficient (CPCC) (Zeng et al. 2022) to OT-CPCC by introducing EMD, which is essentially equivalent to assign different importance to each pair of data points when computing the gradients.
+ It presents a linear time approximation algorithm to solve the OP problem for learning the hierachical representation.
+ Extensive experiments are provided.
summary: The paper extends the work (Zeng et al. 2022) for learning structured representation by introducing the Earth Mover's Distance (EMD) to measure the pairwise distances among classes in embeddings. Moreover, a so-called Fast FlowTree algorithm is proposed to improve the computation efficiency, which is linear in the size of the dataset. Experiments show promising results.
weaknesses: 1.  To replace the $\ell_2$ distance between two class-centroids, the EMD is introduced based on representing class label by all samples in the class. If the number of samples in each class is large, then solving a single EMD is expensive in time cost. Notice that the distances are computed for pairwise classes, thus the times to compute EMD is quadratic to the number of classes. If the number of classes is large, then it is extremely expensive to compute the EMD based hierarchy information among embeddings.  While different approximation methods can be used, it is not clear about the computation cost, compared to the baseline methods. 
 
2. In Table 3, the preciseness of hierachical information listed show that in most cases, the presented FastFT based method cannot beat the $\ell_2$ norm based method. How to interprete the results in Table 3?

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: see Weaknesses
soundness: 3
strengths: The proposed algorithms are novel and scalable. The proposed method leads to performance improvements across multiple datasets. The work is impactful for applications needing hierarchical class embedding.
summary: This paper presents a method for learning structured representations by embedding class hierarchies using Optimal Transport (OT) in the Cophenetic Correlation Coefficient (CPCC) framework. Traditional methods use Euclidean distances between class means to represent hierarchical information, which can be inaccurate for multi-modal distributions. The authors propose using Earth Mover’s Distance (EMD) to address this limitation and introduce Fast FlowTree (FastFT), a computationally efficient OT approximation algorithm for CPCC that operates in linear time.
weaknesses: 1. I feel section 3.3 can be simplified by moving some derivations and technical discussion to the appendix. 
2. Label embedding methods have been employed in extreme classification (see datasets: http://manikvarma.org/downloads/XC/XMLRepository.html). The authors may want to experiment on those datasets to verify scalability and compare to existing embedding methods.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: While the paper proposes FastFT for computational efficiency, can the authors provide a clear comparison of computation time on empirical datasets?
soundness: 3
strengths: 1. The paper is well-written and easy to follow.
2. The paper makes a theoretical extension by generalizing CPCC with EMD, allowing the model to handle more complex, non-unimodal distributions that Euclidean-based methods struggle with. This extension is theoretically sound and shows how OT can capture nuanced hierarchical structures in data.
summary: The paper addresses the challenge of embedding hierarchical structures in feature representations for supervised classification tasks. It builds upon existing work that uses the Cophenetic Correlation Coefficient (CPCC) to enforce hierarchical structure alignment but innovates by employing the Earth Mover’s Distance (EMD) within this framework. The authors also propose a new, efficient algorithm, Fast FlowTree (FastFT), to approximate EMD in linear time, making it feasible to apply OT-based regularization to large datasets.
weaknesses: 1. The improvement of FastFT and EMD, which appears to be the paper’s key contribution, is modest on empirical datasets, showing less than a 1% gain on average across three random seeds.
2. The computation time for empirical datasets is not reported.
3. The paper is a follow-up to the CPCC work by Zeng et al. (2022), with the addition of EMD. The novelty of introducing Fast FlowTree is somewhat limited.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: # Suggestions

- In Sec. 2.2, I suggest the authors to comment more specifically on the complexity of OT and its variants. I get that this is somewhat done in Table 1, but I did not find any references to this table. I think a clear discussion on the text would be beneficial (and, of course, a reference to the table).

- For better uniformity and clarity, the authors might want to change references to the EMD as a distance in favor of $2-$Wasserstein distance, which is what is actually used in the paper. For instance, in most tables the authors use EMD and terms involving the Wasserstein distance (TWD or SWD).

- In section 3.3, when authors are comparing the gradients of $\ell_{2}-$CPCC and OT-CPCC, I think the authors can make 2 further links, if they agree with the following discussion. First, the comparison of distributions' centroids can be understood as computing the MMD, with a linear kernel, between those distributions. This can further make the link with other probability metrics, and arguably, OT yields a richer geometry than this simplistic choice. Second, if one chooses to approximate OT through the Sinkhorn algorithm, letting the entropic regularization penalty $\epsilon \rightarrow +\infty$ leads to the trivial coupling $P_{ij} = 1/nm$. In this case, OT-CPCC boils down to the $\ell_{2}-$CPCC algorithm, modulo a scaling factor on the gradients. The authors can also cite (Feydy et al, 2019) which discusses these relations in depth.

# Questions

I have only one question, which I detail below.

Proposition 3.1 left me wondering if it would not be possible to run OT-CPCC with a parametric approximation of the classes through Gaussian measures. This way, instead of calculating EMD between classes, one would use the Bures-Wasserstein metric,

$\mathcal{W}_{2}(P,Q)^{2}  = \lVert \mu\_{P} - \mu\_{Q} \rVert\_{2}^{2}  +  Tr(\Sigma\_{P} + \Sigma\_{Q} - 2(\Sigma\_{P}^{1/2}\Sigma\_{Q}\Sigma\_{P}^{1/2})^{1/2}).$

With this approach, you could trade the complexity with respect the number of samples to a complexity with respect number of dimensions, which may be desirable depending on the problem at hand. Some authors have used this in other contexts (e.g., (Alvarez-Melis and Fusi, 2020)).

Furthermore, stressing the idea that class-conditionals are multi-modal, GMMs could be fit on these conditionals, then compute Mixture-Wasserstein (Delon and Desolneux, 2020) distances between the class-conditionals.
soundness: 4
strengths: - The ideia is good. OT makes a good contribution to the considered problem, especially concerning how distributions are compared and embedded in the tree structure.

- The authors present a solid analysis of how to extend the $\ell_{2}$-CPCC framework to OT, especially concerning differentiating through linear programming.

- The authors' experiments show a solid improvement over previous framework under a variety of ways of computing Wasserstein distances.
summary: This paper tackles the problem of structured knowledge embedding in supervised learning, i.e., how to enforce that representations encode structures in the label space. An example of structure in label space is when labels are hierarchical (e.g., dog -> husky, bulldog, etc). Previous work has considered comparing class-conditional distributions through the Euclidean distance of the mean of their samples. However, as the authors remark, this choice is restrictive as it does not really capture all the information of a distribution (e.g., multi-modality, covariances, etc). The authors propose using the EMD to compare these class-conditionals, improving the learning of representations.

---

__Post-Rebuttal__

The authors did a good job on the rebuttal, and addressed all of my concerns. I also praise the initiative of running further experiments with the Bures-Wasserstein metric, which proved to be too computationally heavy for the task at hand (since $d \gg n$).

Overall, this is a strong submission, so I decided to increase my score from __6. Marginally above the acceptance threshold__ to __8. Accept, good paper__.
weaknesses: - In Table 1 (Sec 2.2), I do not think that these complexities are straightforward to derive. Can the authors point to the reasoning to their derivation? For instance, EMD's complexity is tipically $\mathcal{O}(n^{3}\log n)$ instead of $\mathcal{O}(n(d+n^{2}\log n))$

- While I got what the authors are trying to express with prop. 3.1, and the result is true if stated correctly, I think the authors could rephrase it for better clarity. First, "$\ell_{2}$ for Gaussians" is not mathematically correct, since the $\ell_{2}$ norm or distance is defined for vectors. The closest term to this expression would be an $L_{2}$ norm or distance between probability distributions, but it is not equivalent to the $\ell_{2}$ distance between the means. I propose that the authors change the statement to "EMD reduces to $\ell_{2}$ distance between the means". Other ways of rephrasing this are possible, see my suggestions and questions below. Likewise, an appropriate reference to (Takatsu, 2011) would be interesting.

- In table 4, the objective "Flat" is not explained. From the overall context, I guess this corresponds to learning with no structure tree, but this is not clear to me.

Meta-review: $\ell_2$-CPCC has been used to embed structured knowledge within labels into feature representations, but has the limitation that it may misrepresent class relationships, especially in multi-modal cases. To address this limitation, the paper propose a EMD-CPCC that use earth mover's distance instead of using class means. Since EMD has high computational cost, the paper further proposes Fast FlowTree which is a linear-time approximation method. Experiments show the benefits of the proposed methods.

The reviewers point out that the paper is well-written and easy to follow, the proposed method is novel and scalable, the experiments are extensive and demonstrates the improvement, and the paper has strong theoretical analysis.

The main concerns of the reviewers were:
- Theory/technical details: proofs are vague, can be simplified, or can be clarified further
- Novelty is lacking: this is a follow-up to the CPCC work by Zeng et al. (2022)
- EMD-based computation is expensive when we have many samples in each class or when we have many classes
- Experiments: FastFT sometimes fails to outperform the baseline or gains are modest
- Suggestions for more contents: literature review on label representation learning, experiments on extreme classification datasets, results on computation time

The authors have updated the paper to address most of the concerns directly, or they have added discussions about them in the paper. 

Overall, the reviewers feel positive and the average rating of this paper became higher after the discussions. I recommend accepting this paper to the conference.

Decision: Accept (Poster)