Source-linked AI summary
hyperbolix: Hyperbolic Deep Learning in JAX
Timo Klein, Thomas Lang, Yllka Velaj, Sebastian Tschiatschek
TL;DR
Hierarchical data motivates hyperbolic methods, but JAX users lack a comprehensive, general-purpose hyperbolic deep-learning library. hyperbolix provides a common manifold interface and broad layer families for JAX and Flax NNX, with independently checked operations. Its cancellation-free hyperboloid formulas retain float32 accuracy at distances where geoopt returns NaN, while reported measurements do not establish training quality on published tasks.
Problem
JAX has Riemannian optimizers but no general-purpose hyperbolic layer library, while hierarchical data is poorly represented by flat vector spaces.
Source: §1 · p. 1Method
hyperbolix provides six manifolds with a common interface, broad neural-network layer families, and supporting optimization and dimensionality-reduction tools for JAX and Flax NNX.
Source: §1 · p. 2Results
Takeaways & Limitations
The library offers JAX users a shared interface for hyperbolic models and layers, with checked operations and improved hyperboloid numerical accuracy.
Source: §1 · p. 2Reported measurements are single-seed sweeps, and no training-quality benchmarks establish published-task accuracy.
Source: §7 · p. 11
Abstract
from arXiv · showhide
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\'e ball, the hyperboloid, the $\kappa$-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 .
1 Introduction
Hierarchical data is poorly matched to flat vector spaces, while JAX users lack a general-purpose hyperbolic layer library. hyperbolix fills that gap and reports independently checked operations and improved hyperboloid numerical accuracy.
1 Introduction
- JAX users otherwise lack a general-purpose hyperbolic layer library and must implement layers from scratch or hand-port them, a process prone to numerical-stability issues.Hierarchical data has tree structure that flat vector spaces cannot represent efficiently; hyperbolic volume expands exponentially with radius, matching tree growth.Source: §1 · pp. 1–2
- hyperbolix provides six manifolds, 76 layer implementations across major layer families, Riemannian optimizers, wrapped distributions, and dimensionality-reduction methods.The library is built for JAX and Flax NNX.Source: §1 · p. 2
- Every checked operation is compared with an independently derived reference rather than another library call.The references are described in Section 5.Source: §1 · p. 2
- In float32, hyperbolix’s cancellation-free hyperboloid formulas stay at rounding-level precision to geodesic radius 10, where geoopt returns NaN.On the Poincaré ball, it matches or exceeds geoopt and HypLL in round-trip precision.Source: §1 · p. 2
2 Design and API
hyperbolix keeps manifold operations compatible with JAX transformations through pure, unbatched operations, stateless manifolds, and call-time curvature. The API also supports learnable curvature, Euclidean and Riemannian optimization, and initialization designed for hyperbolic layers.
2.1 Core design principles
- Pure, unbatched manifold operations and stateless manifold instances make JAX transformations composable while leaving batching layouts to callers.Curvature is passed at call time, so the same manifold instance can serve fixed or learned curvature and be shared across layers.Source: §2.1 · p. 3
2.3 Optimization
- The example converts a Euclidean MLP by projecting features with exp0 under jax.vmap, using hyperbolic layers and a regression head, while leaving the model class and training loop unchanged.The hyperbolic layer can accept separate input and output curvatures; the models use plain optax.adam because their weights are Euclidean parameters.Source: §2.1 · p. 3, Figure 1 · p. 4, §2.3 · p. 4
2.2 Learnable curvature
- LearnableCurvature stores a scalar parameter and offers log, softplus, or identity parameterizations, with optional clamping.The identity parameterization permits signed curvature for the κ-stereographic manifold; an unconstrained parameter lets plain optax.adam update it.Source: §2.2 · p. 3
2.3 Optimization
- Riemannian Adam and SGD are provided as Optax transformations that detect manifold-valued parameters; most models instead use optax.adam for Euclidean parameters.Riemannian optimization applies to parameters that genuinely live on the manifold, such as embedding tables or legacy Ganea-style biases.Source: §2.3 · pp. 3–4
2.4 Initialization
- Initialization must avoid weights that are too large, which can send outputs into numerically unstable regions, or too small, which can collapse distinct inputs to indistinguishable outputs.The small-weight failure compounds a per-layer gain below one until the network becomes a constant function.Source: §2.4 · p. 4
3 The hyperbolix Library
hyperbolix builds its library around a shared manifold interface and offers multiple layer implementations with differing trade-offs. It also includes supporting packages and compares its coverage with existing geometry and deep-learning libraries.
3.1 Manifolds
- The six-manifold library provides a common 19-operation interface, enabling user-defined manifolds to work with manifold-generic components without subclassing.The models differ in optimization and numerical properties; exact isometries connect the three hyperbolic models, and signed curvature lets κ-stereographic geometry cross the Euclidean point smoothly.Source: §3.1 · p. 5
3.2 Neural network layers
- The layer package includes linear, convolution, attention, normalization, positional-encoding, regression, activation, and vector-quantization implementations, with variants that trade off fidelity, numerical stability, and speed.For example, Ganea et al. (2018) use manifold biases requiring Riemannian optimization, while the HNN++ reformulation (Shimizu et al., 2021) keeps parameters Euclidean for optax.adam.Source: §3.2 · pp. 5–6
3.3 Beyond layers
- Beyond layers, hyperbolix provides Riemannian optimizers, wrapped and uniform distributions, HoroPCA and CO-SNE, and Gromov δ-hyperbolicity estimation.Its documentation includes layer-selection tables and a numerical-stability guide, including advice on when to use float64.Source: §3.3 · p. 6
3.4 Comparison to existing libraries
- Compared with existing libraries, hyperbolix is presented as the first comprehensive, general-purpose hyperbolic deep-learning library for JAX; the cited JAX alternatives do not offer a general-purpose hyperbolic layer library.The comparison also includes PyTorch libraries geoopt and HypLL, and Geomstats, which covers general Riemannian geometry but does not focus on deep learning.Source: §3.4 · p. 6
4 Numerical Stability
Hyperboloid operations can lose precision when their standard formulas subtract large, nearly equal terms; hyperbolix rewrites these calculations to avoid that cancellation. Tests show improved stability across the measured ranges, alongside distinct precision limits for the Poincaré ball and hyperboloid.
Failure mode
- Standard hyperboloid formulas lose precision because the Minkowski inner product subtracts nearly equal terms that grow as e^2a.The resulting absolute error can exceed the distance being measured; in float32, cancellation exhausts the mantissa around radius 9, while stored coordinates remain accurate to about radius 16.Source: §4 · pp. 6–7
Cancellation-free formulas
- Hyperbolix avoids ambient cancellation by rewriting Minkowski quadratic forms as sums of non-negative terms, separating radial and angular gaps.The same radial/perpendicular decomposition is applied to tangent metrics, parallel transport, and gyrovector addition.Source: §4 · p. 7
Numerical results
- For curvature c=1 in dimension 8, hyperbolix’s float32 distance between same-ray hyperboloid points one unit apart stays within 8× input-rounding error beyond geodesic radius a≈10, and its round-trip error remains at rounding level for r=10−3 to 20.At r=10−3 near the origin, geoopt’s median hyperboloid round-trip error is 2.3 × 10−2, versus 1.1 × 10−7 for hyperbolix.Source: §4 · p. 7, Figure 2 · p. 8
Round-trip comparisons
- The Poincaré-ball expmap-logmap round trip is more stable for hyperbolix and HypLL than geoopt near the origin; HypLL has no hyperboloid operations.On the hyperboloid, hyperbolix stays at rounding error across the measured tangent norms, while geoopt has elevated error near the origin.Source: §4 · p. 7, Figure 3 · p. 8
Model-specific stability
- For origin-anchored operations, the Poincaré ball loses accuracy earlier than the hyperboloid and proper-velocity models.At curvature c = 1, the ball's float32 projection margin caps geodesic distance at 12.6; the hyperboloid and proper-velocity models remain near rounding error through distance 20.Source: §4 · pp. 6, 9
5 Correctness: Testing Against Independent Oracles
The authors test checked operations against independently derived references rather than relying on agreement between implementations. Their example shows why this matters: a shared reference-code mistake can pass generic manifold checks.
Oracle-based verification
- Every checked operation is tested against an independent oracle: a derived closed form, a NumPy/SciPy transcription, or a finite-difference estimate.The suite contains 1,038 test functions across 55 files, parametrized over float32 and float64.Source: §5 · p. 10
Scope of verification
- A transcription oracle checks implementation against the source paper, not whether the paper's formula itself is correct.The authors therefore distinguish oracle comparisons from consistency checks such as jit versus eager execution and expmap-logmap round trips.Source: §5 · p. 10
Illustrative bug
- In Hyperboloid.log_radius_concat, swapping digamma arguments makes concatenated radius O(N) too large while still producing a valid manifold point.The same swap appears in the “Intrinsic Lorentz Neural Network” paper and reference implementation (Shi et al., 2026); generic on-manifold tests would not catch it.Source: §5 · p. 10
- The oracle caught the digamma-argument swap using the identity E[log ∥v∥] = 1/N, rather than reproducing the reference implementation's mistake.Table 20 reports per-operation deviations between hyperbolix calls and their oracles over standard sampling grids.Source: §5 · p. 10
6 Performance
The performance section explains how JAX compilation and vectorization reduce the overhead of pointwise manifold operations, then compares compiled primitives against other libraries. HyperLL is fastest for its Poincaré-ball implementation, while geoopt and hyperbolix have function-dependent results.
Execution overhead
- Manifold operations chain memory-bound pointwise steps, and eager execution adds a kernel launch and memory round trip for each step.These operations can include norms, clamps, conformal rescalings, and re-projections.Source: §6 · p. 11
JAX execution
- jax.jit fuses many steps and reduces Python dispatch, while jax.vmap batches the single-point manifold definitions across batch shapes.The large benchmark batch is used to amortize dispatch overhead; geoopt and HypLL use torch.compile to match jax.jit.Source: §6 · p. 11
Benchmark results
- HypLL is fastest on the Poincaré ball; geoopt and hyperbolix are broadly comparable, with the faster library depending on the function.The comparison measures distance, exp0, log0, Möbius addition, and parallel transport on CUDA in float32 at batch size 107.Source: §6 · p. 11
7 Limitations
The reported measurements are narrow in scope, and the library’s downstream evidence remains limited. Some numerical and feature-coverage limits are structural.
7 Limitations
- The measurements are single-seed sweeps, and the paper reports no training-quality benchmarks on published tasks.Speed-run brackets bound launch-to-launch machine and compilation state, not sampling variation.Source: §7 · p. 11
- The library was developed and hardened within one research project, with two vector-quantization bottlenecks ported from it.Wider use is needed to expose interfaces that only appear general.Source: §7 · p. 11
- Hyperbolic metric growth imposes a float32 range ceiling, and cancellation-free hyperboloid formulas push the failure point outward without removing it.At c = 1, float64 is the only Poincaré-ball remedy past about radius 14, and it fails past about 24; attention and positional encodings are hyperboloid-only, while vector-quantization bottlenecks are Poincaré-only.Source: §7 · p. 11
8 Conclusion
The conclusion presents hyperbolix as a general-purpose hyperbolic deep-learning library for JAX and Flax NNX, with six manifolds, broad layer families, and independently verified checked operations. The layer inventory details the scope and implementation choices of those components.
8 Conclusion
- Hyperbolix provides six manifolds, broad layer families, optax-based Riemannian optimizers, and independent oracle tests; the appendix enumerates 76 public layer exports across ten families.Its API is designed for jit, vmap, and scan, and cancellation-free hyperboloid arithmetic keeps float32 distance error within 8× input-rounding error.Source: §8 · p. 11, A Layer Inventory · p. 16
A.1 Linear
- Only HypLinearPoincare and HypRegressionPoincare have manifold-valued weights requiring a Riemannian optimizer; other listed linear layers use optax.adam.The library includes 11 linear-layer exports.Source: Table 3 · p. 16
A.2 Convolution
- LorentzConv2D is retained to reproduce the LResNet convolutional stack, while HypConv2DHyperboloid supersedes it for new work.LorentzConv2D uses Euclidean convolution on spatial components and reconstructs the time coordinate from the hyperboloid constraint.Source: Table 4 · p. 17
A.3 Attention
- The attention family includes O(N) kernel-trick linear attention and two O(N^2) alternatives operating in the spatial domain or on full Lorentzian inner products.The linear variant uses the focus function of focus_transform.Source: Table 5 · p. 17
A.4 Normalization & Dropout
- Normalization includes intrinsic gyrovector normalizers and HRC wrappers around Euclidean Flax NNX layers; frechet_variance is manifold-generic.The table also lists dropout and a shared function.Source: Table 6 · p. 17
A.5 Residual, Pooling & Primitives
- The residual, pooling, and primitive family includes Lorentz centroids for weighted midpoints and Einstein/gyromidpoints for Poincaré batch normalization and VQ codebook updates.lorentz_midpoint supports full-attention aggregation and hyperboloid batch-mean estimation.Source: Table 7 · p. 18
A.6 Positional Encoding
- Positional encoding offers learnable relative encoding through HypformerPositionalEncoding and deterministic, rotation-based HyperbolicRoPE.The latter is also named hope.Source: Table 8 · p. 18
A.7 Regression / MLR; A.8 Vector Quantization; A.9 Activations
- The remaining listed families include Euclidean-logit classification heads, explicit EMA and implicit MLR vector-quantization codebooks, and activations applied in the tangent space at the origin.The explicit codebook uses a copy-gradient straight-through estimator; the implicit codebook uses Gumbel-Softmax quantization-as-classification.Source: Table 9 · p. 19, Table 10 · p. 19, Table 11 · p. 19
A.10 Hybrid / Interface
HyperPPFeatureScaling prepares Euclidean features before their first mapping onto a hyperbolic manifold.
A.10 Hybrid / Interface
- HyperPPFeatureScaling applies parameter-free RMSNorm, a Lipschitz activation, and 1/√d scaling before expmap_0, optionally rescaling features to keep the first on-manifold point inside the boundary.The optional learned rescale is bounded by ρmax = atanh(α)/√c.Source: Table 12 · p. 20
B Implemented Methods
The appendix maps published methods to hyperbolix components, surveys comparable JAX libraries, and specifies a shared manifold protocol. It also documents exact isometries among the shipped point-set models.
B Implemented Methods
- The inventory also lists implementations for HNN, HNN++, and Poincaré ResNet, including Poincaré linear, convolutional, and normalization components.HNN++ is identified as point-to-hyperplane FC/MLR.Source: B Implemented Methods · p. 21
- The implemented-method inventory maps Fully Hyperbolic Neural Networks to hyperboloid layers and activations, and fully hyperbolic CNNs to their listed components.The appendix’s table maps 31 published-method rows to hyperbolix components.Source: B Implemented Methods · p. 21, Table 13 · p. 21
C The JAX Ecosystem
- Among the surveyed JAX libraries, none offers a general-purpose hyperbolic layer library; GeoJAX has Poincaré and hyperboloid models and broad Riemannian optimizers, but no neural-network layers.Rieoptax has nine undocumented Poincaré/Lorentz layer classes and no convolutions; RiemannAX’s hyperbolic manifolds are in unreleased development code.Source: C The JAX Ecosystem · p. 22, Table 14 · p. 22
D The Manifold Protocol
- The manifold protocol uses structural typing for a common 19-method interface, and the three shipped point-set models have exact distance-preserving isometries in both directions.The hyperboloid–proper-velocity mapping is direct rather than composed through the Poincaré ball.Source: D The Manifold Protocol · pp. 22–23, Table 15 · p. 23, Table 16 · p. 23
E Initialization Defaults
The shipped initialization defaults target per-layer input-Jacobian gain near 1, while numerical conventions address gradient safety, tolerance limits, precision choices, and test coverage.
E Initialization Defaults
- Near-1 input-Jacobian gain guides shipped layer-family initialization, with fan-in or fan-out scaling used where source-paper initialization is not reproduced.Overshooting can push outputs toward manifold boundaries and cause NaN losses; undershooting can compound across layers until distances quantize to zero and gradients vanish.Source: E Initialization Defaults · pp. 23–24, Table 17 · p. 24
F Numerical Conventions
- MIN_NORM floors norms and denominators, while input and output tanh clamps keep downstream atanh away from its singularity.The norm floor gives safe_norm a finite gradient at zero; output clipping addresses float32 tanh saturation that input clipping alone does not prevent.Source: F Numerical Conventions · p. 25
- default_atol uses √ε for dimensionless manifold-constraint residuals, but its near-origin calibration can reject valid distant points; float64 is recommended in that regime.The stated reliable ranges are approximately distance 10 in float64 and 7 in float32.Source: F Numerical Conventions · p. 25, Table 18 · p. 25
- Manifold-operation compute dtype and layer-parameter storage dtype are independent; float64 compute with float32 storage is the recommended high-precision recipe.The recipe preserves geometric compute precision while using half the parameter and optimizer-state memory.Source: F Numerical Conventions · pp. 25–26
G Test-Suite Structure
- At v1.3.0, pytest parametrization expands 1,038 top-level test functions into 4,446 collected test items.The suite spans dtype, curvature, and layer-spec axes.Source: G Test-Suite Structure · p. 26
G.1 The 18-property layer contract
The appendix checks layer contracts and independently verifies numerical accuracy, stability, and runtime across hyperbolix and available baselines. The results show strong accuracy and stability, with speed varying by operation, manifold, and runtime setup.
G.1 The 18-property layer contract
- The 18-property layer contract checks output shape, dtype, finiteness, manifold membership, eager-versus-jit agreement, curvature cases, gradients, tangent-input behavior, and convolution arithmetic.The contract includes finite, nonzero gradients and requires tangent-input results to agree with explicit expmap_0 followed by the manifold-input branch.Source: §G.1 · p. 26
G.2 Oracle taxonomy; G.3 Wrapped-normal log-probability: four-direction pinning
- Every numeric assertion is checked against one of three independent oracle sources, never another hyperbolic-geometry library; wrapped-normal log_prob is pinned for both models from four independent directions.At c = 10^-4, both models must reduce to a Euclidean Gaussian in Riemannian coordinates, matching scipy.stats.multivariate_normal.logpdf to 10^-3.Source: §G.2 · p. 27, §G.3 · p. 27
H.1 Common protocol
- Speed comparisons distinguish compiled kernels from eager dispatch, aggregate five launches, and treat structurally unavailable cells as absent rather than unmeasured.The CPU one-core control is the like-for-like comparison; default-threading results reflect different runtime thread defaults.Source: §H.1 · p. 28, §H.6 · pp. 33–34
H.2 Deviation from independent oracles
- Every float32 oracle-deviation cell is below 4.6 × 10^-6 and every float64 cell is at or below 3.5 × 10^-14, within the stated membership tolerances.The deviations compare library operations with independent NumPy float64 oracles over sampled curvature and dimension grids.Source: §H.2 · p. 28, §H.1 · p. 28, Table 20 · p. 29
H.3 Cross-library accuracy on the Poincaré ball
- In float32, the three Poincaré implementations differ by at most 2.2× per operation-curvature-dimension cell; in float64, curvature storage limits HypLL and the natural geoopt call at c = 0.3.hyperbolix passes curvature in the working dtype at call time and has no corresponding precision floor.Source: §H.3 · p. 29, Table 21 · p. 30
H.4 Round-trip stability
- On the hyperboloid, hyperbolix stays below the float32 breakdown threshold through radius 20, while geoopt exceeds it at radius 10^-3 with median relative error 2.3 × 10^-2.On the Poincaré ball, the breakdown radius is 7.24 for hyperbolix versus 3.38 for geoopt and HypLL.Source: §H.4 · p. 30
H.5 Speed on an NVIDIA H100
- At batch 107, HypLL leads three Poincaré primitives, while hyperbolix is 1.9–2.9× faster than geoopt on three hyperboloid primitives.On the Poincaré ball, Möbius addition and parallel transport are also roughly 2.5× slower for hyperbolix than the PyTorch libraries, consistent with XLA splitting reductions from rescaling. HypLL has no hyperboloid model, geoopt has no layers, and geoopt's compiled distance and logmap0 fall back to host dispatch.Source: §H.1 · p. 28, §H.5 · p. 31, Table 22 · p. 32
H.6 Speed on a CPU
- At one core and batch 105, hyperbolix is fastest on four of nine primitive rows, including disjoint five-launch brackets for Poincaré log0 and hyperboloid exp0 and log0.The fourth win, Poincaré exp0, is a tie with HypLL because the brackets overlap; most default-threading gaps shrink under the one-core control.Source: §H.6 · p. 34, Table 24 · p. 35
H.7 Two-point cancellation on the hyperboloid
This section tests cancellation in hyperboloid two-point operations with exact constructions, comparing hyperbolix and geoopt across radii and precisions.
H.7 Two-point cancellation on the hyperboloid
- The distance error grows like e^2a in both libraries because the Minkowski inner product subtracts two O(e^2a) terms away from the origin.The two points used for the distance construction lie on the same ray at radii a and a + 1, so their exact distance is 1.Source: §H.7 · p. 35
- Hyperbolix’s float32 distance error reaches 7.2 × 10^-4 at a = 14 and 1.0 × 10^-1 at a = 16, so its stability is not unconditional.Above a = 8, geoopt’s maximum error is reduction-order noise; its median is the meaningful comparison.Source: §H.7 · p. 36
- Geoopt’s float32 distance returns NaN from a = 9, while hyperbolix stays at float32 rounding level through a = 10 and within 6.7 × 10^-16 in float64 through a = 16.At a = 12, hyperbolix’s float32 maximum absolute deviation is 3.7 × 10^-4 (median 5.6 × 10^-5).Source: Table 26 · p. 36, §H.7 · p. 36
- For tangent norms, hyperbolix stays within 1.35× the input-rounding error from a = 6 in the tested hyperboloid setup.The test uses unit-norm tangent vectors directed 45° from the ray and compares both libraries against the exact norm.Source: §H.7 · p. 35, Table 26 · p. 36