Hyperbolix offers versatile hyperbolic deep learning tools in jax
hyperbolix: Hyperbolic Deep Learning in JAX
Machine LearningMathematical Software
Summary
Deep learning often happens in flat (Euclidean) spaces, but some problems benefit from curved spaces called hyperbolic spaces. The authors created hyperbolix, a collection of tools in the JAX programming framework to make working with different types of hyperbolic spaces easier and more reliable. It supports many mathematical models and neural network layers, plus special optimizations to keep calculations precise where others fail. This lets developers build advanced neural networks that better capture complex relationships in data with curved geometry.
What this means in practice
- •For machine learning engineers: Develop neural networks using hyperbolic geometry for tasks that model hierarchical or complex relational data more effectively.
- •For software developers in scientific computing: Integrate hyperbolic space operations with JAX for high-precision and efficient numerical computations involving curved geometries.
Authors
Timo Klein, Thomas Lang, Yllka Velaj, Sebastian Tschiatschek
Abstract
We present hyperbolix, an open-source library for hyperbolic deep learning in JAX, built on Flax NNX. To our knowledge, it is the first comprehensive, general-purpose hyperbolic deep learning library in JAX. It includes six manifolds with a common interface: Euclidean space, the Poincaré ball, the hyperboloid, the $κ$-stereographic model, mixed-curvature product spaces, and the proper velocity space. We implement layer families that cover linear layers, convolutions, attention, normalization, positional encoding, regression, and vector quantization. These building blocks span methods ranging from Ganea's original hyperbolic neural networks to recent fully hyperbolic architectures such as Hypformer and Lorentzian ResNet. Additionally, hyperbolix contains Riemannian optimizers implemented as optax transformations, wrapped distributions, and hyperbolic dimensionality-reduction techniques. Its API uses idiomatic JAX: Manifolds are stateless, with curvature being passed at call time, while manifold operations act on single points, with jax.vmap enabling batch operations. The precision of every checked operation is tested against a closed-form NumPy/SciPy transcription from the source paper or a finite difference, for both float32 and float64. On the hyperboloid, standard formulas for two-point operations, such as the distance, lose precision far from the origin, because they subtract two large, nearly equal terms. hyperbolix replaces these subtractions with cancellation-free formulas that stay accurate in float32 at distances where prior implementations return NaN. hyperbolix is available under the MIT license at https://github.com/timoklein/hyperbolix .