Joint Hierarchical Representation Learning of Samples and Features via Informed Tree-Wasserstein Distance
Ya-Wei Eileen Lin, Ronald R. Coifman, Gal Mishne, Ronen Talmon
TL;DR
This work tackles the challenge of learning meaningful hierarchical representations in high-dimensional data where both samples and features exhibit structure. It introduces an unsupervised, iterative framework that alternates learning trees for one data mode and computing Tree-Wasserstein Distances (TWD) for the other, thereby refining both hierarchies. A data-driven Haar wavelet filtering step, constructed from the learned trees, further enhances the representations, with theoretical convergence guarantees. The learned hierarchies can initialize hyperbolic graph networks, yielding improved link prediction and node classification, and empirical results on word-document and scRNA-seq datasets show superior performance in sparse approximation and unsupervised metric learning tasks. Overall, the method provides a principled, scalable approach to jointly uncovering and leveraging hierarchical structure across both rows and columns of high-dimensional data, with practical impact for downstream hyperbolic architectures.
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.