Representation Learning of Geometric Trees

TL;DR

This work tackles representation learning for geometric trees, where geometry, topology, and hierarchical order jointly shape structure. It introduces GTMP, a branch-focused message passing scheme that is invariant to SE(3) transformations and preserves geometric-topological information with linear time complexity, along with GT-SSL, a self-supervised framework that exploits hierarchical ordering and subtree growth via frequency-domain geometry and Earth Mover’s Distance. The combination yields strong discriminative representations, demonstrated across eight real-world datasets (neurons and rivers) with substantial gains over baselines, and shows robust transferability and invariance properties. Together, GTMP and GT-SSL provide scalable, geometry-aware learning for complex tree-structured data, enabling improved downstream tasks in domains like neuroscience and geomorphology.

Abstract

Geometric trees are characterized by their tree-structured layout and spatially constrained nodes and edges, which significantly impacts their topological attributes. This inherent hierarchical structure plays a crucial role in domains such as neuron morphology and river geomorphology, but traditional graph representation methods often overlook these specific characteristics of tree structures. To address this, we introduce a new representation learning framework tailored for geometric trees. It first features a unique message passing neural network, which is both provably geometrical structure-recoverable and rotation-translation invariant. To address the data label scarcity issue, our approach also includes two innovative training targets that reflect the hierarchical ordering and geometric structure of these geometric trees. This enables fully self-supervised learning without explicit labels. We validate our method's effectiveness on eight real-world datasets, demonstrating its capability to represent geometric trees.

Figures (5)

• Figure 1: (a) Illustration of a neuron's geometric tree-like structure; (b) Representation of a river network exhibiting a tree structure embedded within a geometric landscape; (c) Three different geometric trees with isomorphic network connectivity and identical spatial coordinates. Distinguishing these geometric trees requires jointly considering spatial, topology, and hierarchical layout information.
• Figure 2: Illustration of the GTMP model. The geometric information is first extracted on each length-three branch starting from node $v_i$, then they are aggregated with other node information to update the node embedding $\mathbf{h}_i$.
• Figure 3: Illustration of the hierarchical relationship between tree nodes through partial ordering. In this scenario, $T_k$ is a subtree of $T_j$, establishing a partial ordering relationship between their respective subtree embeddings. Conversely, since $T_{k^{\prime}}$ does not constitute a subtree of $T_j$, there is no requirement for a partial ordering relationship between the embeddings of these two subtrees.
• Figure 4: Robustness test on rotation and translation invariance by augmenting the data on the test set. The x-axis corresponds to the magnitude of the rotation angle (left) and translation distance (right) while y-axis shows the AUC scores. We can observe that our proposed GTMP model stays invariant to both translation and rotation transformations.
• Figure 5: Images of example neural cell reconstruction from NeuroMorpho.org, which depict the reconstruction for the cell with NeuroMorpho.org ID NMO_86952 ascoli2007neuromorphokoch2016big. The left image is a screenshot of the cell as viewed in the “Animation” feature on the cell’s corresponding NeuroMorpho.org page, and the right image is the cell’s representative image in the database.