Hierarchical Mamba Meets Hyperbolic Geometry: A New Paradigm for Structured Language Embeddings

TL;DR

This paper introduces Hierarchical Mamba (HiM), a novel framework combining hyperbolic geometry with sequence models to improve hierarchical language understanding and reasoning in NLP tasks.

Contribution

The paper proposes HiM, integrating Mamba2 with hyperbolic geometry, enabling hierarchy-aware embeddings that outperform Euclidean models on multiple hierarchical NLP tasks.

Findings

HiM effectively captures hierarchical relationships in language data.

HiM variants outperform Euclidean baselines on four datasets.

HiM-Poincaré provides detailed hierarchical distinctions, HiM-Lorentz offers robustness.

Abstract

Selective state-space models excel at long-sequence modeling, but their capacity for language representation -- in complex hierarchical reasoning -- remains underexplored. Most large language models rely on \textit{flat} Euclidean embeddings, limiting their ability to capture latent hierarchies. To address this, we propose {\it Hierarchical Mamba (HiM)}, integrating efficient Mamba2 with hyperbolic geometry to learn hierarchy-aware language embeddings for deeper linguistic understanding. Mamba2-processed sequences are projected to the Poincar\'e ball or Lorentzian manifold with ``learnable'' curvature, optimized with a hyperbolic loss. Our HiM model facilitates the capture of relational distances across varying hierarchical levels, enabling effective long-range reasoning for tasks like mixed-hop prediction and multi-hop inference in hierarchical classification. Experimental results show both HiM variants effectively capture hierarchical relationships across four linguistic and medical datasets, surpassing Euclidean baselines, with HiM-Poincar\'e providing fine-grained distinctions with higher h-norms, while HiM-Lorentz offers more stable, compact, and hierarchy-preserving embeddings-favoring robustness. The source code is publicly available at https://github.com/BerryByte/HiM.