Joint Hierarchical Representation Learning of Samples and Features via Informed Tree-Wasserstein Distance

TL;DR

This paper introduces an unsupervised joint hierarchical representation learning method for high-dimensional data with hierarchical structures in samples and features, using Tree-Wasserstein Distance to iteratively refine both hierarchies.

Contribution

It proposes a novel alternating approach to jointly learn hierarchical trees for samples and features via Tree-Wasserstein Distance, with theoretical convergence guarantees.

Findings

Improves hyperbolic graph convolutional networks for link prediction and node classification.

Outperforms baselines in sparse approximation tasks.

Effective in unsupervised Wasserstein distance learning on biological and text datasets.

Abstract

High-dimensional data often exhibit hierarchical structures in both modes: samples and features. Yet, most existing approaches for hierarchical representation learning consider only one mode at a time. In this work, we propose an unsupervised method for jointly learning hierarchical representations of samples and features via Tree-Wasserstein Distance (TWD). Our method alternates between the two data modes. It first constructs a tree for one mode, then computes a TWD for the other mode based on that tree, and finally uses the resulting TWD to build the second mode's tree. By repeatedly alternating through these steps, the method gradually refines both trees and the corresponding TWDs, capturing meaningful hierarchical representations of the data. We provide a theoretical analysis showing that our method converges. We show that our method can be integrated into hyperbolic graph convolutional networks as a pre-processing technique, improving performance in link prediction and node classification tasks. In addition, our method outperforms baselines in sparse approximation and unsupervised Wasserstein distance learning tasks on word-document and single-cell RNA-sequencing datasets.