well_poincare_rl

Hierarchical multi-step field predictor in the Poincaré ball, with optional PPO fine-tuning, continual learning (Replay + EWC), and explicit data/checkpoint contracts for multi-contributor scientific use.

This model predicts spatiotemporal scientific fields and embeds hierarchical structure (e.g. taxonomies) using hyperbolic geometry.

Model description

Component Role
MultiScaleEncoder Channel-agnostic encoder (shared 1×1 stem per channel → mean fusion → spectral + local path) → 8-D Euclidean latent
HierarchicalHyperbolicPredictor Multi-step prediction in the Poincaré ball (coarse RNN + residual refinement levels)
HierarchyEmbedding Poincaré vs Euclidean node embeddings for taxonomy trees (RiemannianAdam / geoopt)
Continual stack Per-domain FieldNormalizer, same-C ReplayBuffer, DiagonalEWC, optional hyperbolic distillation
Optional RL Gymnasium env + PPO / on-policy fine-tune with hyperbolic critic

Intended uses

  • Multi-step forecasting of continuum fields (reaction–diffusion, active matter, fluids, similar PDE-style data)
  • Continual training across scientific domains with different channel counts (no fixed-C encoder lock-in)
  • Hierarchy recovery on taxonomic / tree-structured scientific data (explicit Poincaré vs Euclidean comparison)
  • Research on hyperbolic representations for multi-scale and hierarchical scientific structure

How to use

pip install torch geoopt gymnasium h5py the_well
# clone / install this package, then:
import torch
from src.model import MultiScaleEncoder, HierarchicalHyperbolicPredictor
from src.config import BEST_HPARAMS as BEST

enc = MultiScaleEncoder(hidden=BEST["hidden"], out_dim=8)
model = HierarchicalHyperbolicPredictor(
    enc, c=BEST["curvature"], pred_steps=BEST["pred_steps"], levels=BEST["levels"]
)
# x: (B, T, C, H, W) or (B, C, H, W) — C is not fixed by the encoder
pred = model(x)  # (B, pred_steps, 8) points in the Poincaré ball

Real multi-stream continual (requires HF network access):

python -m src.run_multistream \
  --datasets gray_scott_reaction_diffusion active_matter shear_flow \
  --max-samples 96 --epochs-per-domain 3

Hierarchy embedding (real PBDB or synthetic tree):

python -m src.run_hierarchy_embed \
  --pbdb-taxa Dinosauria Mammalia \
  --dims 8 --loss-types softmax --burn-in-epochs 0 \
  --c-values 1.0 2.0 --epochs 80 --lr 0.02 --seeds 0 --optimizer radam

Synthetic-only smoke tests: python -m src.run_full --synthetic and python -m src.run_hierarchy_embed --synthetic ....

Training data

Source Role Provenance label
The Well (HF streams) Spatiotemporal fields (e.g. gray_scott, active_matter, shear_flow) REAL_STREAMED
Local Well-format HDF5 Same, offline REAL_LOCAL
PBDB occurrence records Taxonomy edges (Dinosauria + Mammalia) and optional density fields REAL_PBDB / REAL_PBDB_TAXONOMY
Synthetic Well-like / synthetic trees Opt-in only (--synthetic); never silent fallback SYNTHETIC / SYNTHETIC_TREE

Contract: real-data paths hard-fail on missing data, schema mismatch, or stream failure. Synthetic data is only available via explicit APIs. Provenance is always reported and stored with checkpoints.

Evaluation results

Verification tiers used below:

  • 🔒 Checkpoint-verified — recomputed from a saved artifact (hash chain, stored metrics, or model structure)
  • 📋 Reported — from a run log; consistent and plausible, not independently re-derived from a saved artifact

Track A — Multi-stream continual learning (Well)

Sequential training on live HF streams, no synthetic fallback.

🔒 Verified from checkpoint: REAL_STREAMED provenance, dataset names, per-domain normalizer channel counts (C=2, 11, 4), channel-agnostic encoder architecture.

📋 Reported retention losses (checkpoint format does not yet persist the loss history):

After domain gray_scott (C=2) active_matter (C=11) shear_flow (C=4)
Domain 1 0.3486
Domain 2 0.3561 0.3299
Domain 3 0.3248 0.3484 0.3358

Replay buffer by channel count after full run: {2: 24, 11: 24, 4: 24}.

Track B — Real PBDB hierarchy (Poincaré vs Euclidean)

Live taxonomy edges from Dinosauria + Mammalia (~16k occurrence records → 2141 nodes, 2172 edges).

🔒 Independently recomputed from raw occurrence JSON using this repo’s build_edge_list / hash_edge_list / hash_config / combined_identity_hash; matched checkpoint identity and stored results_table to full precision.

Setting ΔMRR (Poincaré − Euclidean)
dim=8, softmax, c=1.0 🔒 +0.1151
dim=8, softmax, c=2.0 🔒 +0.1576

Protocol: reconstruction (train = test edges) — measures embedding capacity at a given dimension, not link-prediction generalization.

Scope: one seed, one dimension, metrics-only hierarchy checkpoint (full embedding weights not exported in the verified artifact).

Hyperparameters (default / Optuna best)

lr=3.82e-4, curvature=0.455, hidden=96, batch_size=8,
pred_steps=4, w_phys=9.6e-4, levels=2, window=4

Single source of truth: src/config.py (BEST_HPARAMS).

Limitations

  • Multistream retention numbers are 📋 reported, not 🔒 checkpoint-verified (loss history not stored in the multistream checkpoint format).
  • Real PBDB hierarchy result is one seed / one dimension.
  • Channel fusion is mean-pooling (lossy); attention-based fusion is future work.
  • Soft hierarchical region structure in the Poincaré ball is currently deferred.
  • While PPO is implemented and smoke-tested; it has not yet been optimized and verified for performance.
  • Conservation auxiliary loss in the shipped physics suite is a soft proxy; a stricter spatial-integral form is a recommended upgrade, not necessarily present in every checkpoint.
  • Supported compute: CPU and CUDA. Not ported to TPU/XLA or LPU.
  • Data licenses are separate from code: PBDB is CC0; The Well has its own terms. This card’s Apache-2.0 license covers code and released model artifacts, not third-party datasets.

Ethical considerations

  • In many fields, when a measure becomes a target it ceases to be a good measure; that is likely true with reward hacking this design as well.
  • This release is for research and methodological development.
  • The field of Paleobiology includes data with known incompleteness and occasional inconsistencies;
  • As models trained on this system reflect curator-assigned taxonomy in PBDB, they reflect such inconsistencies across recorded embeddings.
  • The pipeline records edges as-observed and does not “resolve” taxonomic disputes.

Citation

If you use this code or results, please cite the repository and the relevant data sources (The Well, PBDB). A formal paper citation will be added when the arXiv draft is public.

@software{well_poincare_rl,
  title = {well\_poincare\_rl: Hierarchical hyperbolic field prediction and taxonomy embedding},
  year = {2026},
  license = {Apache-2.0}
}

License

Code and released model artifacts: Apache License 2.0.

Dataset terms remain those of the original providers (The Well, Paleobiology Database, etc.).

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support