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-
Cencoder 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 likley 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 inconsisistencies 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.).