Add HI-Mapper source (code only, no weights)

.gitattributes

@@ -1,35 +1,5 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,92 @@
+---
+license: apache-2.0
+library_name: pytorch
+tags:
+  - hyperbolic
+  - lorentz
+  - hierarchy
+  - entailment
+  - pedestrian-attribute-recognition
+  - clip
+  - promptpar
+  - hi-mapper
+pipeline_tag: feature-extraction
+---
+
+# HI-Mapper (code only)
+
+Hyperbolic hierarchy mapper for pedestrian attribute recognition (PromptPAR).
+This Hub repo contains **source code only** — **no trained weights**.
+
+HI-Mapper lifts Euclidean region / CLIP features into the Lorentz hyperboloid,
+builds a fixed depth-3 anatomical tree, and supervises it with entailment cones,
+sibling separation, and radius ordering (MERU / HyCoCLIP style). An optional
+attribute-grounded entailment term maps PETA-style attribute prefixes onto tree
+nodes. An optional HypDAE-style hyperbolic diffusion decoder is included.
+
+## Contents
+
+```
+hi_mapper/
+  __init__.py          # public exports
+  lorentz.py           # Lorentz manifold + EuclideanToLorentz lift
+  tree.py              # hierarchical / attribute entailment losses
+  hi_mapper.py         # DivHiMapper + PETA attr grouping
+  hyp_diffusion.py     # optional hyperbolic diffusion decoder
+  ARCHITECTURE.md      # detailed architecture notes
+```
+
+## Install
+
+```bash
+pip install torch
+# then copy hi_mapper/ into your project, or:
+git clone https://huggingface.co/ZACK777/hi-mapper
+```
+
+Requires PyTorch. No Hub weights are downloaded.
+
+## Quick start
+
+```python
+import torch
+from hi_mapper import DivHiMapper, build_attr_groups
+
+# region tokens: [B, 5, D] = global + 4 anatomical leaves
+B, D = 2, 768
+region_tokens = torch.randn(B, 5, D)
+cls = torch.randn(B, D)
+
+mapper = DivHiMapper(
+    feat_dim=D,
+    curvature=0.2,
+    target_radius=1.0,
+    attr_groups=None,  # or build_attr_groups(attr_names)
+)
+root, mid, leaves, hier_loss, prompt_loss, attr_loss = mapper(
+    region_tokens, cls
+)
+```
+
+## Geometry notes
+
+- `EuclideanToLorentz` uses a running-norm scaler (not LayerNorm) so relative
+  norms — and thus hyperbolic radii — stay meaningful.
+- Distance uses a numerically stable `asinh` form; Minkowski products run in
+  float64 intermediates.
+- Entailment uses half-aperture / exterior-angle cones (`K=0.1`).
+
+## Integration
+
+Designed as a drop-in branch for PromptPAR / CLIP-based PAR models. See
+`hi_mapper/ARCHITECTURE.md` for the tree layout, losses, and optional decoder.
+
+## Citation / context
+
+Built for hyperbolic hierarchical regularisation on top of PromptPAR-style
+region tokens. Related ideas: MERU, HyCoCLIP, HypDAE, Khrulkov et al. (δ-hyperbolicity).
+
+## License
+
+Apache-2.0 (code). Upstream PromptPAR / CLIP licenses still apply if you use
+those backbones.

example.py

@@ -0,0 +1,10 @@
+"""Minimal smoke test for HI-Mapper (no weights)."""
+import torch
+from hi_mapper import DivHiMapper
+
+B, D = 2, 768
+mapper = DivHiMapper(feat_dim=D, curvature=0.2, target_radius=1.0)
+root, mid, leaves, hier_loss, prompt_loss, attr_loss = mapper(
+    torch.randn(B, 5, D), torch.randn(B, D)
+)
+print("root", tuple(root.shape), "hier_loss", float(hier_loss))

hi_mapper/ARCHITECTURE.md

@@ -0,0 +1,340 @@
+# Hi-Mapper + Hyperbolic Generative Decoder
+
+Detailed architecture of the **Hi-Mapper** plug-in module: hyperbolic hierarchy mapping and an optional HypDAE-style generative decoder.
+
+---
+
+## Overview
+
+The module takes a set of **region-level visual tokens** and builds a **depth-3 binary tree** in Lorentz hyperbolic space. It produces:
+
+1. A **refined root feature** for downstream use.
+2. A **hierarchical contrastive loss** that enforces parent–child and sibling relations between levels.
+3. Optionally, a **prompt-hierarchy alignment loss** when auxiliary prompt tokens are provided.
+4. Optionally, a **diffusion reconstruction loss** that treats the hierarchy as a generative code for a global feature (train only).
+
+```
+region_tokens [B, 5, D]     global token + 4 region leaves
+global_feat   [B, D]        optional anchor (e.g. CLS)
+prompt_nodes  [vis_depth, …] optional, for prompt alignment
+        │
+        ▼
+┌─────────────────────────────┐
+│         DivHiMapper         │
+│  Euclidean tree → Lorentz   │
+│  hierarchical contrastive   │
+│  optional prompt alignment  │
+└─────────────────────────────┘
+        │
+        ├── root_feat [B, D]
+        ├── mid_feats [B, 2, D]
+        ├── leaf_feats [B, 4, D]
+        ├── hier_loss
+        └── prompt_loss (optional)
+                │
+                ▼
+┌─────────────────────────────┐
+│  HyperbolicDiffusionDecoder   │  (train only, optional)
+│  z = hierarchy code → CLS     │
+└─────────────────────────────┘
+        │
+        └── diff_loss
+```
+
+**Combined module loss:**
+
+```
+hi_mapper_loss = hier_loss
+               + prompt_hier_w * prompt_loss     # optional
+               + hyp_diffusion_w * diff_loss     # optional
+```
+
+---
+
+## 1. Inputs and outputs
+
+### DivHiMapper (`hi_mapper.py`)
+
+| Tensor | Shape | Role |
+|--------|-------|------|
+| `region_tokens` | `[B, 5, D]` | Index 0 = global region token; indices 1–4 = four anatomical leaves (head, upper, lower, feet) |
+| `global_feat` | `[B, D]` | Optional global anchor (soft-combined into root) |
+| `prompt_deep` | `[L, P, 1, W]` | Optional; pooled into 5 prompt nodes for alignment |
+
+| Output | Shape | Role |
+|--------|-------|------|
+| `root_feat` | `[B, D]` | Hierarchy-refined global representation |
+| `mid_feats` | `[B, 2, D]` | Upper-half and lower-half mid nodes |
+| `leaf_feats` | `[B, 4, D]` | Leaf node features (unchanged from input leaves) |
+| `hier_loss` | scalar | Lorentz hierarchical contrastive loss |
+| `prompt_loss` | scalar | Prompt tree + alignment loss (0 if disabled) |
+
+Default feature dimension `D = 768` (CLIP ViT-L/14 projected dim).
+
+---
+
+## 2. Tree topology (depth 3)
+
+Fixed binary tree — no learned tree search (anatomy is predefined):
+
+```
+                         root  (whole entity)
+                        /                    \
+                 mid_upper                 mid_lower
+                /         \               /          \
+            leaf_0       leaf_1       leaf_2       leaf_3
+           (head)       (upper)       (lower)      (feet)
+
+  region_tokens[:, 0] (global) ──► soft-anchors root via global_gate
+  global_feat         ──► soft-anchors root via (1 - global_gate)
+```
+
+Parent table (`tree.py`):
+
+- Leaves 0,1 → mid 0; leaves 2,3 → mid 1.
+- Mids 0,1 → root.
+
+---
+
+## 3. Euclidean tree construction
+
+### Leaves
+
+```python
+global_tok = region_tokens[:, 0]      # [B, D]
+leaves     = region_tokens[:, 1:5]    # [B, 4, D]
+```
+
+### Mid nodes — PairMerge
+
+Two sibling leaves are merged with a small MLP:
+
+```
+PairMerge(a, b):
+  concat(a, b) → Linear(2D → D) → GELU → Linear(D → D)
+```
+
+```python
+mid_upper = PairMerge(leaves[:, 0], leaves[:, 1])
+mid_lower = PairMerge(leaves[:, 2], leaves[:, 3])
+mid       = stack([mid_upper, mid_lower], dim=1)   # [B, 2, D]
+```
+
+### Root
+
+```python
+root_raw = PairMerge(mid[:, 0], mid[:, 1])
+root     = root_raw + gate * global_tok + (1 - gate) * global_feat
+```
+
+`gate` is a learnable scalar (`global_gate`, init 0.5).
+
+---
+
+## 4. CLIP Euclidean → Lorentz conversion block
+
+**Code:** `lorentz.py` — `EuclideanToLorentz`, `expmap0_euclidean_space`
+
+CLIP features (`all_class`, CLS, merged tree nodes) live in **Euclidean** space ℝ^D. Hierarchical losses need Lorentz points on the hyperboloid. Tree merges (`PairMerge`) stay Euclidean; conversion happens **after** the tree is built.
+
+```
+CLIP / tree nodes (Euclidean ℝ^D)
+        │
+        ▼
+┌─────────────────────────────────┐
+│     EuclideanToLorentz            │
+│  Linear(D→D) adapter              │
+│  α · v   (α learnable, init 1/��D) │
+│  MERU expmap0_euclidean_space     │
+└─────────────────────────────────┘
+        │
+        ▼
+Lorentz points [time, space] ∈ ℝ^{D+1}
+```
+
+### MERU-style lift (not Minkowski-norm on CLIP vectors)
+
+Treat Euclidean embedding `v ∈ ℝ^D` as **space components only** (tangent at the hyperboloid origin):
+
+```
+r        = clamp(√c ‖v‖₂, max_norm)
+x_space  = unit(v) · sinh(r) / √c
+x_time   = sqrt(1/c + ‖x_space‖²)
+h        = [x_time, x_space] ∈ ℝ^{D+1}
+```
+
+Learnable scale `α` (init `1/√D`) is applied before the expmap to keep `sinh` numerically stable for CLIP-scale norms (MERU, ICML 2023).
+
+### LorentzManifold
+
+- Curvature `c` (default `1.0`, flag `--hi_mapper_curvature`).
+- **`from_euclidean` / `EuclideanToLorentz`:** correct CLIP lift (above).
+- **`geodesic_dist(x, y)`:** Lorentz distance via Minkowski inner product and `acosh` (on already-lifted points).
+
+All tree nodes (leaves, mid, root) are mapped independently:
+
+```
+leaves_h  [B, 4, D+1]
+mid_h     [B, 2, D+1]
+root_h    [B, 1, D+1]
+```
+
+Hyperbolic space is used because tree volume grows exponentially with depth; Euclidean embeddings distort those relations. Downstream classification still uses the **Euclidean** `root_feat`.
+
+---
+
+## 5. Hierarchical contrastive loss
+
+**Code:** `tree.py` — `hierarchical_contrastive_loss`
+
+Margin-based loss (`margin = 0.1`) on geodesic distances in Lorentz space:
+
+| Constraint | Meaning |
+|------------|---------|
+| `d(leaf, parent_mid) < d(leaf, uncle_mid)` | Each leaf pulled toward its parent, not the sibling branch |
+| `d(mid, root) < d(mid, sibling_mid)` | Mid nodes organized under root |
+| `d(leaf, mid) < d(leaf, root)` | Leaves shallower than root (depth consistency) |
+
+Formally, for each constraint:
+
+```
+loss += ReLU(d_parent - d_other + margin).mean()
+```
+
+This implements the Hi-Mapper idea: **child–parent similar, siblings/uncles dissimilar**, in a shared hyperbolic space.
+
+---
+
+## 6. Prompt-hierarchy alignment (optional)
+
+**Flag:** `--optimize_prompts_with_hi_mapper`  
+**Weight:** `--prompt_hier_w` (default `0.05`)
+
+When `prompt_deep` is provided:
+
+1. **Pool** prompt tensor into 5 groups (same grouping as region tokens): mean over depth and token dims → `[5, D]`.
+2. **Project** to feature dim if needed (`prompt_proj`: `W → D`).
+3. **Build the same Lorentz tree** on prompt nodes.
+4. **Losses:**
+   - `prompt_hier`: same hierarchical contrastive loss on the prompt tree.
+   - `align`: geodesic distance between prompt Lorentz nodes and **batch-mean** of visual Lorentz nodes (visual side detached).
+
+```
+prompt_loss = prompt_hier + align
+```
+
+This gives a direct training signal so auxiliary prompt parameters follow the same hierarchy as the visual region tokens.
+
+---
+
+## 7. HyperbolicDiffusionDecoder (generative aspect)
+
+**Code:** `hyp_diffusion.py`  
+Inspired by **HypDAE** (ICCV 2025): the hierarchy acts as a semantic code; a small denoiser learns to reconstruct a global feature from that code.
+
+### Hierarchy code
+
+```python
+z = concat(root_feat, mid_feats.flatten(1), leaf_feats.flatten(1))
+# [B, D * (1 + 2 + 4)] = [B, 5376]  when D=768
+```
+
+### DDPM step (feature space)
+
+Training only — **not used at inference**.
+
+1. Sample `t ~ Uniform{0, …, steps−1}`, `ε ~ N(0, I)`.
+2. Noisy target: `x_t = √ᾱ_t · global_feat + √(1−ᾱ_t) · ε`.
+3. Predict noise: `ε̂ = eps_θ(x_t, z, t_emb)`.
+4. Loss: `MSE(ε̂, ε)`.
+
+### Network
+
+```
+t_emb = MLP(t / steps)                    # [B, 128]
+inp   = [x_t ‖ z ‖ t_emb]                 # [B, D + 5376 + 128]
+ε̂     = Linear → SiLU → Linear → SiLU → Linear
+```
+
+Default `steps = 6` (`--hyp_diffusion_steps`).
+
+The decoder forces the hierarchy to be **generatively sufficient** for the global feature, without pixel-level image synthesis.
+
+---
+
+## 8. Module map
+
+| File | Class / symbol | Role |
+|------|----------------|------|
+| `lorentz.py` | `EuclideanToLorentz`, `expmap0_euclidean_space`, `LorentzManifold` | CLIP Euclidean→Lorentz lift; geodesic distance; curvature |
+| `tree.py` | `DIV_TREE`, `hierarchical_contrastive_loss`, `alignment_loss` | Fixed parent table and losses |
+| `hi_mapper.py` | `DivHiMapper`, `PairMerge` | Euclidean tree build; Lorentz lift; prompt alignment |
+| `hyp_diffusion.py` | `HyperbolicDiffusionDecoder` | CLS / global feature reconstruction DDPM |
+
+---
+
+## 9. Hyperparameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `feat_dim` (`D`) | 768 | Feature dimension |
+| `div_num` | 4 | Number of region leaves (5 tokens total with global) |
+| `hi_mapper_curvature` | 1.0 | Lorentz curvature `c` |
+| `hi_mapper_w` | 0.1 | Weight on `hi_mapper_loss` in total training loss |
+| `prompt_hier_w` | 0.05 | Weight on prompt alignment branch |
+| `hyp_diffusion_w` | 0.1 | Weight on diffusion MSE |
+| `hyp_diffusion_steps` | 6 | DDPM timesteps |
+
+---
+
+## 10. Training vs inference
+
+| Component | Training | Inference |
+|-----------|----------|-----------|
+| DivHiMapper (tree + Lorentz) | ✓ | ✓ |
+| `root_feat` output | ✓ | ✓ |
+| Hierarchical contrastive loss | ✓ | — |
+| Prompt alignment | optional | — |
+| HyperbolicDiffusionDecoder | optional (aux loss) | **off** |
+
+At inference, only the forward pass through DivHiMapper runs; no diffusion sampling, no extra loss terms.
+
+---
+
+## 11. Data flow (tensor shapes, D=768)
+
+```
+region_tokens [B,5,768]
+global_feat   [B,768]
+        │
+        ├─ leaves [B,4,768]
+        ├─ mid    [B,2,768]  ← PairMerge pairs
+        └─ root   [B,768]    ← PairMerge + gate blend
+        │
+        ├─ EuclideanToLorentz → leaves_h [B,4,769], mid_h [B,2,769], root_h [B,1,769]
+        │
+        ├─ hier_loss  (scalar)
+        │
+        ├─ prompt_loss (scalar, optional)
+        │
+        └─ z [B,5376] ──► diff_loss (scalar, optional, needs global_feat as target)
+```
+
+---
+
+## 12. Design rationale
+
+- **Fixed tree:** Region structure is known (4 body bands + global); no Mixture-of-Gaussians tree search as in the original Hi-Mapper paper.
+- **Lorentz geometry:** Encodes exponential growth of hierarchy levels with lower distortion than Euclidean space.
+- **Explicit Euclidean→Lorentz block:** CLIP features are Euclidean; MERU-style space-only `expmap0` (with learnable `α`) lifts them correctly onto the hyperboloid without treating the first CLIP coordinate as Minkowski time.
+- **Generative decoder:** HypDAE-style auxiliary loss makes the hierarchy a predictive code for the global feature, improving root representation without generating images.
+- **Prompt alignment:** Optional branch to synchronize learnable prompt parameters with the visual hierarchy in the same manifold.
+
+---
+
+## References
+
+- Kwon et al., *Improving Visual Recognition with Hyperbolical Visual Hierarchy Mapping* (Hi-Mapper), CVPR 2024. [arXiv:2404.00974](https://arxiv.org/abs/2404.00974)
+- Desai et al., *Hyperbolic Image-Text Representations* (MERU), ICML 2023. [arXiv:2304.09172](https://arxiv.org/abs/2304.09172)
+- Li et al., *HypDAE: Hyperbolic Diffusion Autoencoders for Hierarchical Few-shot Image Generation*, ICCV 2025.
+- Code: [kwonjunn01/Hi-Mapper](https://github.com/kwonjunn01/Hi-Mapper)

hi_mapper/__init__.py

@@ -0,0 +1,26 @@
+from .hi_mapper import DivHiMapper, build_attr_groups
+from .hyp_diffusion import HyperbolicDiffusionDecoder
+from .lorentz import (
+    EuclideanToLorentz,
+    LorentzManifold,
+    entailment_loss,
+    exterior_angle,
+    half_aperture,
+)
+from .tree import (
+    attribute_entailment_loss,
+    hierarchical_entailment_loss,
+)
+
+__all__ = [
+    "DivHiMapper",
+    "build_attr_groups",
+    "HyperbolicDiffusionDecoder",
+    "EuclideanToLorentz",
+    "LorentzManifold",
+    "entailment_loss",
+    "exterior_angle",
+    "half_aperture",
+    "attribute_entailment_loss",
+    "hierarchical_entailment_loss",
+]

hi_mapper/hi_mapper.py

@@ -0,0 +1,198 @@
+"""Div-based Hi-Mapper with optional prompt-hierarchy alignment."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from .lorentz import EuclideanToLorentz
+from .tree import alignment_loss, attribute_entailment_loss, hierarchical_entailment_loss
+
+# Node layout used when stacking the tree for attribute grounding:
+#   0..3 leaves (head, upper body, lower body, feet), 4..5 mid, 6 root
+LEAF_HEAD, LEAF_UPPER, LEAF_LOWER, LEAF_FEET = 0, 1, 2, 3
+MID_UPPER, MID_LOWER, ROOT = 4, 5, 6
+
+# PETA attribute-name prefixes mapped onto hierarchy nodes. Whole-person
+# attributes (age, gender) sit at the root because they are the most general;
+# carried objects span the torso, so they attach to the mid nodes; the rest are
+# region-specific and attach to the matching row band.
+_PREFIX_TO_NODE = (
+    ("personal", ROOT),
+    ("carrying", MID_UPPER),
+    ("accessory", LEAF_HEAD),
+    ("hair", LEAF_HEAD),
+    ("head", LEAF_HEAD),
+    ("upperbody", LEAF_UPPER),
+    ("lowerbody", LEAF_LOWER),
+    ("footwear", LEAF_FEET),
+    ("shoes", LEAF_FEET),
+)
+
+
+def build_attr_groups(attr_names) -> dict[int, list[int]]:
+    """
+    Map attribute indices onto hierarchy nodes by name prefix.
+
+    Returns ``{node_index: [attr_index, ...]}``. Attributes whose prefix is not
+    recognised are skipped rather than forced onto a node, so datasets without
+    PETA's naming convention simply produce fewer (or no) grounded constraints.
+    """
+    groups: dict[int, list[int]] = {}
+    for idx, raw in enumerate(attr_names):
+        name = str(raw).strip().lower()
+        for prefix, node in _PREFIX_TO_NODE:
+            if name.startswith(prefix):
+                groups.setdefault(node, []).append(idx)
+                break
+    return groups
+
+
+class PairMerge(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim * 2, dim),
+            nn.GELU(),
+            nn.Linear(dim, dim),
+        )
+
+    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return self.net(torch.cat([a, b], dim=-1))
+
+
+class DivHiMapper(nn.Module):
+    """
+    Build a depth-3 tree from PromptPAR div tokens (all_class) and optional
+    prompt_deep slices, with a Lorentz entailment-cone hierarchy loss.
+
+    Tree merges stay in Euclidean space; CLIP Euclidean features are lifted
+    to the Lorentz hyperboloid via ``EuclideanToLorentz`` (MERU-style) before
+    the hierarchical losses.
+    """
+
+    def __init__(
+        self,
+        feat_dim: int = 768,
+        prompt_width: int = 1024,
+        div_num: int = 4,
+        curvature: float = 1.0,
+        optimize_prompts: bool = False,
+        target_radius: float = 1.0,
+        entail_margin: float = 0.0,
+        attr_groups: dict[int, list[int]] | None = None,
+    ):
+        super().__init__()
+        self.feat_dim = feat_dim
+        self.div_num = div_num
+        self.optimize_prompts = optimize_prompts
+        self.entail_margin = entail_margin
+        self.attr_groups = attr_groups or {}
+
+        self.euclid_to_lorentz = EuclideanToLorentz(
+            feat_dim, curvature=curvature, target_radius=target_radius
+        )
+        # Shared manifold reference for hierarchical / alignment losses
+        self.manifold = self.euclid_to_lorentz.manifold
+
+        # Text embeddings live on a different norm scale than visual tokens, so
+        # they get their own adapter and scale (MERU uses separate lambda_img /
+        # lambda_txt for exactly this reason) while sharing the manifold.
+        self.text_to_lorentz = None
+        if self.attr_groups:
+            self.text_to_lorentz = EuclideanToLorentz(
+                feat_dim, curvature=curvature, target_radius=target_radius
+            )
+            self.text_to_lorentz.manifold = self.manifold
+        self.prompt_proj = nn.Linear(prompt_width, feat_dim) if prompt_width != feat_dim else nn.Identity()
+
+        self.merge_pair = PairMerge(feat_dim)
+        self.merge_root = PairMerge(feat_dim)
+        self.global_gate = nn.Parameter(torch.tensor(0.5))
+        # Zero-init: at step 0 the root is exactly the CLIP CLS token, so
+        # enabling Hi-Mapper cannot damage the classifier before it has learned
+        # anything. Previously a randomly-initialised PairMerge output replaced
+        # CLS outright and cost ~10 mA points in epoch 1.
+        self.root_gate = nn.Parameter(torch.tensor(0.0))
+
+    def _to_lorentz(self, x: torch.Tensor) -> torch.Tensor:
+        """Lift CLIP Euclidean features to Lorentz hyperboloid (MERU-style)."""
+        if x.dim() == 2:
+            x = x.unsqueeze(1)
+        return self.euclid_to_lorentz(x)
+
+    def _build_tree(self, tokens: torch.Tensor, cls_tok: torch.Tensor | None = None):
+        """
+        tokens: [B, div_num+1, D] — index 0 global part token, 1..div_num row bands.
+        Returns euclidean root/mid/leaves and lorentz versions.
+        """
+        assert tokens.shape[1] == self.div_num + 1, (
+            f"Expected {self.div_num + 1} div tokens, got {tokens.shape[1]}"
+        )
+        global_tok = tokens[:, 0]
+        leaves = tokens[:, 1:]
+
+        mid_upper = self.merge_pair(leaves[:, 0], leaves[:, 1])
+        mid_lower = self.merge_pair(leaves[:, 2], leaves[:, 3])
+        mid = torch.stack([mid_upper, mid_lower], dim=1)
+
+        delta = self.merge_root(mid[:, 0], mid[:, 1]) + self.global_gate * global_tok
+        if cls_tok is not None:
+            root = cls_tok + self.root_gate * delta
+        else:
+            root = delta
+
+        leaves_h = self._to_lorentz(leaves)
+        mid_h = self._to_lorentz(mid)
+        root_h = self._to_lorentz(root.unsqueeze(1))
+
+        hier_loss = hierarchical_entailment_loss(
+            leaves_h, mid_h, root_h, self.manifold, margin=self.entail_margin
+        )
+        return root, mid, leaves, leaves_h, mid_h, root_h, hier_loss
+
+    def _pool_prompt_groups(self, prompt_deep: torch.Tensor) -> torch.Tensor:
+        """
+        Pool prompt_deep [vis_depth, prompt_num, 1, width] into [div_num+1, feat_dim].
+        Groups match div slices in clip/model.py (div_prompt_num = prompt_num // (div_num+1)).
+        """
+        vis_depth, prompt_num, _, width = prompt_deep.shape
+        group_size = prompt_num // (self.div_num + 1)
+        nodes = []
+        for g in range(self.div_num + 1):
+            sl = prompt_deep[:, g * group_size : (g + 1) * group_size]
+            pooled = sl.mean(dim=(0, 1, 2))  # [width]
+            nodes.append(self.prompt_proj(pooled))
+        return torch.stack(nodes, dim=0)
+
+    def forward(
+        self,
+        all_class: torch.Tensor,
+        cls_tok: torch.Tensor,
+        prompt_deep: torch.Tensor | None = None,
+        text_features: torch.Tensor | None = None,
+    ):
+        root, mid, leaves, leaves_h, mid_h, root_h, hier_loss = self._build_tree(all_class, cls_tok)
+
+        prompt_loss = torch.zeros((), device=all_class.device, dtype=hier_loss.dtype)
+        if self.optimize_prompts and prompt_deep is not None:
+            prompt_nodes = self._pool_prompt_groups(prompt_deep).unsqueeze(0)  # [1, G, D]
+            prompt_nodes_b = prompt_nodes.expand(all_class.shape[0], -1, -1)
+            _, _, _, p_leaves_h, p_mid_h, p_root_h, prompt_hier = self._build_tree(
+                prompt_nodes_b, cls_tok=None
+            )
+            visual_stack = torch.cat([leaves_h, mid_h, root_h], dim=1)
+            prompt_stack = torch.cat([p_leaves_h, p_mid_h, p_root_h], dim=1)
+            align = alignment_loss(
+                prompt_stack, visual_stack.detach(), self.manifold, margin=self.entail_margin
+            )
+            prompt_loss = prompt_hier + align
+
+        attr_loss = torch.zeros((), device=all_class.device, dtype=hier_loss.dtype)
+        if self.attr_groups and text_features is not None and self.text_to_lorentz is not None:
+            node_stack = torch.cat([leaves_h, mid_h, root_h], dim=1)      # [B, 7, D+1]
+            attr_h = self.text_to_lorentz(text_features.unsqueeze(1)).squeeze(1)  # [A, D+1]
+            attr_loss = attribute_entailment_loss(
+                node_stack, attr_h, self.attr_groups, self.manifold, margin=self.entail_margin
+            )
+
+        return root, mid, leaves, hier_loss, prompt_loss, attr_loss

hi_mapper/hyp_diffusion.py

@@ -0,0 +1,52 @@
+"""HypDAE-style feature-space diffusion decoder (auxiliary loss only)."""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class HyperbolicDiffusionDecoder(nn.Module):
+    """
+    Small DDPM-style denoiser in Euclidean feature space, conditioned on the
+    hierarchy code. Trained to reconstruct the CLS visual feature; not used
+    at inference time.
+    """
+
+    def __init__(self, dim: int, code_dim: int, steps: int = 6, hidden: int = 512):
+        super().__init__()
+        self.dim = dim
+        self.steps = steps
+        betas = torch.linspace(1e-4, 2e-2, steps)
+        alphas = 1.0 - betas
+        alpha_bar = torch.cumprod(alphas, dim=0)
+        self.register_buffer("alpha_bar", alpha_bar)
+
+        t_dim = hidden // 4
+        self.time_mlp = nn.Sequential(
+            nn.Linear(1, t_dim),
+            nn.SiLU(),
+            nn.Linear(t_dim, t_dim),
+        )
+        self.eps_net = nn.Sequential(
+            nn.Linear(dim + code_dim + t_dim, hidden),
+            nn.SiLU(),
+            nn.Linear(hidden, hidden),
+            nn.SiLU(),
+            nn.Linear(hidden, dim),
+        )
+
+    def forward(self, hierarchy_code: torch.Tensor, cls_tok: torch.Tensor) -> torch.Tensor:
+        b = cls_tok.shape[0]
+        device = cls_tok.device
+        t = torch.randint(0, self.steps, (b,), device=device)
+        eps = torch.randn_like(cls_tok)
+        ab = self.alpha_bar[t].view(b, 1)
+        x_t = ab.sqrt() * cls_tok + (1.0 - ab).sqrt() * eps
+
+        t_emb = self.time_mlp((t.float() / max(self.steps - 1, 1)).unsqueeze(-1))
+        inp = torch.cat([x_t, hierarchy_code, t_emb], dim=-1)
+        eps_hat = self.eps_net(inp)
+        return F.mse_loss(eps_hat, eps)

hi_mapper/lorentz.py

@@ -0,0 +1,320 @@
+"""Self-contained Lorentz hyperboloid ops (adapted from Hi-Mapper / MERU, no geoopt)."""
+from __future__ import annotations
+
+import math
+from typing import Union
+
+import torch
+import torch.nn as nn
+
+# Tangent radius cap. Only guards against overflow: sinh(8) ~ 1.5e3 is safe in
+# fp32. Anything near 1.0 silently discards feature magnitude and zeroes the
+# gradient flowing through the radius.
+EXP_MAX_NORM = 8.0
+
+# Absolute bounds for the learnable curvature (log space).
+MIN_CURVATURE = 0.01
+MAX_CURVATURE = 10.0
+
+
+def _eps(dtype: torch.dtype) -> float:
+    """Precision floor appropriate to the dtype (fp64 tests need a tighter one)."""
+    return 1e-15 if dtype == torch.float64 else 1e-7
+
+
+def _sqrt(x: torch.Tensor) -> torch.Tensor:
+    return torch.sqrt(torch.clamp_min(x, _eps(x.dtype) ** 2))
+
+
+def acosh(x: torch.Tensor) -> torch.Tensor:
+    e = _eps(x.dtype)
+    x = torch.clamp_min(x, 1.0 + e)
+    z = torch.sqrt(torch.clamp_min(x.pow(2) - 1.0, e))
+    return torch.log(x + z)
+
+
+def _inner(u: torch.Tensor, v: torch.Tensor, keepdim: bool = False, dim: int = -1) -> torch.Tensor:
+    """
+    Minkowski inner product on ambient Lorentz coords [time, space].
+
+    Always evaluated in float64 and returned in float64. The expression
+    ``-u_t·v_t + <u_s, v_s>`` cancels catastrophically once the hyperbolic
+    radius grows (the time component scales as cosh(r), so at r=7 the two terms
+    are ~1e6 apart from a result of order 1). In fp32 that destroys every digit
+    by r≈5; the surrounding ops are cheap at these tensor sizes, so the extra
+    precision costs effectively nothing.
+    """
+    d = u.size(dim) - 1
+    uv = u.double() * v.double()
+    if not keepdim:
+        return -uv.narrow(dim, 0, 1).squeeze(dim) + uv.narrow(dim, 1, d).sum(dim=dim, keepdim=False)
+    return -uv.narrow(dim, 0, 1) + uv.narrow(dim, 1, d).sum(dim=dim, keepdim=True)
+
+
+def _norm(u: torch.Tensor, keepdim: bool = False, dim: int = -1) -> torch.Tensor:
+    return _sqrt(_inner(u, u, keepdim=keepdim, dim=dim))
+
+
+def project(x: torch.Tensor, k: Union[float, torch.Tensor], dim: int = -1, max_norm: float = EXP_MAX_NORM) -> torch.Tensor:
+    if not torch.is_tensor(k):
+        k = torch.tensor(k, device=x.device, dtype=x.dtype)
+    dn = x.size(dim) - 1
+    right = x.narrow(dim, 1, dn)
+    if max_norm:
+        right = torch.renorm(right, 2, dim if dim >= 0 else dim - x.dim(), max_norm)
+    left = _sqrt((1.0 / k) + (right * right).sum(dim=dim, keepdim=True))
+    return torch.cat((left, right), dim=dim)
+
+
+def expmap0_euclidean_space(
+    v: torch.Tensor,
+    k: Union[float, torch.Tensor],
+    dim: int = -1,
+    max_norm: float = EXP_MAX_NORM,
+) -> torch.Tensor:
+    """
+    MERU-style expmap0 for Euclidean space vectors.
+
+    Treats ``v ∈ ℝ^D`` as space components only (tangent at the hyperboloid
+    origin). Does NOT apply Minkowski norm to CLIP Euclidean features.
+
+      x_space = sinh(√c ‖v‖) / (√c ‖v‖) · v
+      x_time  = sqrt(1/c + ‖x_space‖²)
+
+    Returns Lorentz coords ``[..., D+1]`` with time first: ``[x_time, x_space]``.
+    """
+    if not torch.is_tensor(k):
+        k = torch.tensor(k, device=v.device, dtype=v.dtype)
+    # Euclidean L2 norm of space components
+    v_norm = torch.linalg.vector_norm(v, dim=dim, keepdim=True).clamp_min(1e-8)
+    # Clamp tangent radius r = √c ‖v‖ for numerical stability (MERU / Hi-Mapper)
+    radius = (v_norm * _sqrt(k)).clamp_max(max_norm)
+    # x_space = sinh(√c ‖v‖) / (√c ‖v‖) · v  ≡  unit(v) · sinh(r) / √c
+    v_unit = v / v_norm
+    x_space = v_unit * (torch.sinh(radius) / _sqrt(k))
+    x_time = _sqrt((1.0 / k) + (x_space * x_space).sum(dim=dim, keepdim=True))
+    return torch.cat((x_time, x_space), dim=dim)
+
+
+def expmap0(u: torch.Tensor, k: Union[float, torch.Tensor], dim: int = -1) -> torch.Tensor:
+    """
+    Exponential map at origin for ambient Lorentz tangent vectors [time, space].
+
+    Prefer ``expmap0_euclidean_space`` when lifting CLIP Euclidean features.
+    """
+    if not torch.is_tensor(k):
+        k = torch.tensor(k, device=u.device, dtype=u.dtype)
+    nomin = _norm(u, keepdim=True, dim=dim).to(u.dtype)
+    safe = nomin.clamp_min(1e-8)
+    u_unit = u / safe
+    nomin = nomin.clamp_max(EXP_MAX_NORM)
+    l_v = torch.cosh(nomin)
+    r_v = torch.sinh(nomin) * u_unit
+    dn = r_v.size(dim) - 1
+    return torch.cat((l_v + r_v.narrow(dim, 0, 1), r_v.narrow(dim, 1, dn)), dim=dim)
+
+
+def dist(x: torch.Tensor, y: torch.Tensor, k: Union[float, torch.Tensor], dim: int = -1) -> torch.Tensor:
+    """
+    Lorentz geodesic distance, evaluated in the numerically stable form
+
+        d(x, y) = (2/√c) · asinh( √( c · ⟨x-y, x-y⟩_L / 4 ) )
+
+    which is algebraically identical to the textbook ``acosh(-c⟨x,y⟩_L)/√c``.
+
+    The textbook form is unusable for nearby points: ``-c⟨x,y⟩_L`` approaches 1
+    and ``acosh(1+ε) ≈ √(2ε)`` amplifies the square root of the coordinate
+    error, so fp32 inputs give d(x, x) ≈ 2e-3 instead of 0. Differencing first
+    avoids the cancellation entirely: when x == y the difference is exactly
+    zero, so the distance is exactly zero.
+    """
+    if not torch.is_tensor(k):
+        k = torch.tensor(k, device=x.device, dtype=x.dtype)
+    diff = x.double() - y.double()
+    sq = _inner(diff, diff, dim=dim, keepdim=False).clamp_min(0.0)
+    kd = k.double()
+    return (2.0 * torch.asinh(_sqrt(kd * sq / 4.0)) / _sqrt(kd)).to(x.dtype)
+
+
+def half_aperture(x: torch.Tensor, k: Union[float, torch.Tensor], min_radius: float = 0.1) -> torch.Tensor:
+    """
+    Half-aperture of the entailment cone at ``x`` (MERU eq. 10, arbitrary curvature).
+
+        aper(x) = asin( 2K / (√c ‖x_space‖) ),  K = min_radius = 0.1
+
+    The cone narrows as ``x`` moves away from the origin, so points near the
+    origin (general concepts) can entail many children.
+    """
+    if not torch.is_tensor(k):
+        k = torch.tensor(k, device=x.device, dtype=x.dtype)
+    x_space_norm = torch.linalg.vector_norm(x[..., 1:], dim=-1).clamp_min(1e-8)
+    ratio = 2.0 * min_radius / (_sqrt(k) * x_space_norm)
+    return torch.asin(ratio.clamp(-1.0 + 1e-6, 1.0 - 1e-6))
+
+
+def exterior_angle(x: torch.Tensor, y: torch.Tensor, k: Union[float, torch.Tensor]) -> torch.Tensor:
+    """
+    Exterior angle ∠Oxy between the cone axis at ``x`` and the geodesic to ``y``
+    (MERU eq. 11, arbitrary curvature).
+    """
+    if not torch.is_tensor(k):
+        k = torch.tensor(k, device=x.device, dtype=x.dtype)
+    cxy = k.double() * _inner(x, y, dim=-1, keepdim=False)
+    x_space_norm = torch.linalg.vector_norm(x[..., 1:].double(), dim=-1).clamp_min(1e-12)
+    num = y[..., 0].double() + x[..., 0].double() * cxy
+    den = x_space_norm * _sqrt(cxy.pow(2) - 1.0)
+    ratio = (num / den.clamp_min(1e-12)).clamp(-1.0 + 1e-9, 1.0 - 1e-9)
+    return torch.acos(ratio).to(x.dtype)
+
+
+def entailment_loss(
+    parent: torch.Tensor,
+    child: torch.Tensor,
+    k: Union[float, torch.Tensor],
+    margin: float = 0.0,
+) -> torch.Tensor:
+    """
+    MERU entailment loss: penalize the child for lying outside the parent's cone.
+
+        L = max(0, ext(parent, child) - aper(parent) + margin)
+
+    ``margin`` (HyCoCLIP, ICLR 2025) pushes children strictly inside the cone
+    rather than merely onto its boundary.
+    """
+    ext = exterior_angle(parent, child, k)
+    aper = half_aperture(parent, k)
+    return torch.relu(ext - aper + margin)
+
+
+class LorentzManifold(nn.Module):
+    """Lorentz manifold with learnable curvature (log-space)."""
+
+    def __init__(self, curvature: float = 1.0, learnable: bool = True):
+        super().__init__()
+        self.log_k = nn.Parameter(torch.tensor(math.log(curvature)), requires_grad=learnable)
+
+    @property
+    def k(self) -> torch.Tensor:
+        return self.log_k.exp()
+
+    def clamp_k(self) -> None:
+        """Clamp curvature to absolute bounds so it cannot drift to 0 or blow up."""
+        with torch.no_grad():
+            self.log_k.data.clamp_(math.log(MIN_CURVATURE), math.log(MAX_CURVATURE))
+
+    def to_hyperboloid(self, tangent: torch.Tensor) -> torch.Tensor:
+        """Lift ambient Lorentz tangent [time, space] via expmap0."""
+        self.clamp_k()
+        return expmap0(tangent, self.k)
+
+    def from_euclidean(self, v: torch.Tensor) -> torch.Tensor:
+        """Lift Euclidean space vectors (CLIP-style) via MERU expmap0."""
+        self.clamp_k()
+        return expmap0_euclidean_space(v, self.k)
+
+    def geodesic_dist(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        self.clamp_k()
+        return dist(x, y, self.k)
+
+    def radius(self, x: torch.Tensor) -> torch.Tensor:
+        """Geodesic distance from the origin, i.e. the point's 'depth'."""
+        self.clamp_k()
+        x_space_norm = torch.linalg.vector_norm(x[..., 1:], dim=-1)
+        return torch.asinh(_sqrt(self.k) * x_space_norm) / _sqrt(self.k)
+
+    def entailment(self, parent: torch.Tensor, child: torch.Tensor, margin: float = 0.0) -> torch.Tensor:
+        self.clamp_k()
+        return entailment_loss(parent, child, self.k, margin=margin)
+
+
+class EuclideanToLorentz(nn.Module):
+    """
+    Explicit CLIP Euclidean → Lorentz hyperboloid conversion (MERU-style).
+
+    Pipeline:
+      v_euclid → Linear(D→D) adapter → v / E[‖v‖] → α · v → expmap0_euclidean_space
+
+    The scale of the tangent vector is the whole ballgame. MERU's ``α = 1/√D``
+    assumes the encoder output has norm ≈ √D, which puts the tangent radius
+    ``√c·α·‖v‖`` at ≈ 1. Raw CLIP ``all_class`` features have norm ≈ 8.3, not
+    √768 ≈ 27.7, and a default-init adapter shrank them a further ≈ 0.58x. The
+    result landed at radius ≈ 0.21, where ``sinh(r) ≈ r`` and the manifold is
+    numerically indistinguishable from Euclidean space.
+
+    The correction divides by a running estimate of the *dataset* mean norm
+    rather than normalising each sample. That distinction matters: per-sample
+    normalisation (LayerNorm, L2) forces every point onto the same radius, and
+    in hyperbolic space the radius *is* the hierarchy signal - general concepts
+    near the origin, specific ones near the boundary. Dividing by a shared
+    scalar fixes the global scale, which was the bug, while leaving the relative
+    magnitudes that encode depth intact.
+
+    ``log_alpha`` is stored in log space (MERU; HyperVLM ICCVW 2025) so the scale
+    cannot collapse to zero or flip sign during training.
+    """
+
+    def __init__(
+        self,
+        feat_dim: int,
+        curvature: float = 1.0,
+        learnable_curvature: bool = True,
+        use_adapter: bool = True,
+        max_norm: float = EXP_MAX_NORM,
+        target_radius: float = 1.0,
+        momentum: float = 0.05,
+    ):
+        super().__init__()
+        self.feat_dim = feat_dim
+        self.max_norm = max_norm
+        self.target_radius = target_radius
+        self.momentum = momentum
+        self.manifold = LorentzManifold(curvature=curvature, learnable=learnable_curvature)
+
+        if use_adapter:
+            self.adapter = nn.Linear(feat_dim, feat_dim)
+            # Identity init: the adapter must not silently rescale the norm.
+            nn.init.eye_(self.adapter.weight)
+            nn.init.zeros_(self.adapter.bias)
+        else:
+            self.adapter = nn.Identity()
+
+        # EMA of the mean feature norm; -1 flags "not yet initialised" so the
+        # first batch seeds it exactly instead of dragging up from a guess.
+        self.register_buffer("running_norm", torch.tensor(-1.0))
+
+        # After the running-norm division the mean norm is 1, so
+        # α = target_radius / √c puts the *mean* radius at `target_radius`.
+        alpha0 = target_radius / math.sqrt(curvature)
+        self.log_alpha = nn.Parameter(torch.tensor(math.log(alpha0)))
+
+    @property
+    def alpha(self) -> torch.Tensor:
+        return self.log_alpha.exp()
+
+    def _scale(self, v: torch.Tensor) -> torch.Tensor:
+        norms = v.norm(dim=-1)
+        if self.training:
+            batch_mean = norms.mean().detach()
+            with torch.no_grad():
+                if self.running_norm.item() < 0:
+                    self.running_norm.fill_(batch_mean.item())
+                else:
+                    self.running_norm.mul_(1.0 - self.momentum).add_(self.momentum * batch_mean)
+        ref = self.running_norm if self.running_norm.item() > 0 else norms.mean().detach()
+        return v / ref.clamp_min(1e-6)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Euclidean features ``[..., D]`` or ``[B, N, D]``.
+        Returns:
+            Lorentz points ``[..., D+1]`` with time first.
+        """
+        *prefix, d = x.shape
+        assert d == self.feat_dim, f"Expected feat_dim={self.feat_dim}, got {d}"
+        # sinh/acosh are unstable in fp16; keep the whole lift in fp32.
+        flat = x.reshape(-1, d).float()
+        v = self.alpha * self._scale(self.adapter(flat))
+        self.manifold.clamp_k()
+        hyp = expmap0_euclidean_space(v, self.manifold.k, max_norm=self.max_norm)
+        return hyp.view(*prefix, d + 1)

hi_mapper/tree.py

@@ -0,0 +1,165 @@
+"""Fixed div hierarchy tree: 4 row-band leaves -> 2 mid -> 1 root."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Tuple
+
+import torch
+import torch.nn.functional as F
+
+from .lorentz import LorentzManifold
+
+
+@dataclass(frozen=True)
+class DivTree:
+    """Parent indices for [leaf0..3, mid0..1, root]."""
+
+    depth: int = 3
+    num_leaves: int = 4
+    num_mid: int = 2
+
+    @property
+    def leaf_to_mid(self) -> Tuple[int, ...]:
+        return (0, 0, 1, 1)
+
+    @property
+    def mid_to_root(self) -> Tuple[int, ...]:
+        return (0, 0)
+
+
+DIV_TREE = DivTree()
+
+
+def hierarchical_entailment_loss(
+    leaves_h: torch.Tensor,
+    mid_h: torch.Tensor,
+    root_h: torch.Tensor,
+    manifold: LorentzManifold,
+    margin: float = 0.0,
+    sibling_margin: float = 0.2,
+    radius_margin: float = 0.1,
+    w_entail: float = 1.0,
+    w_sibling: float = 0.5,
+    w_radius: float = 0.5,
+) -> torch.Tensor:
+    """
+    Hyperbolic hierarchy loss on a depth-3 binary tree (MERU / HyCoCLIP style).
+
+    Three complementary terms:
+
+    1. **Entailment** - each child must lie inside its parent's entailment cone.
+       This is the partial-order constraint that actually encodes a hierarchy in
+       hyperbolic space.
+    2. **Sibling separation** - a leaf stays closer to its own parent than to its
+       uncle, keeping the two branches apart.
+    3. **Radius ordering** - root nearer the origin than mid, mid nearer than
+       leaves, so generality maps onto distance from the origin. Without this the
+       cones are free to sit at arbitrary depths and the hierarchy is unanchored.
+
+    Args:
+        leaves_h: [B, 4, D+1]
+        mid_h: [B, 2, D+1]
+        root_h: [B, 1, D+1]
+    """
+    tree = DIV_TREE
+    root = root_h[:, 0]
+
+    # ---- 1. entailment: parent cone contains child -------------------------
+    entail = []
+    for i in range(tree.num_mid):
+        entail.append(manifold.entailment(root, mid_h[:, i], margin=margin))
+    for i in range(tree.num_leaves):
+        parent = mid_h[:, tree.leaf_to_mid[i]]
+        entail.append(manifold.entailment(parent, leaves_h[:, i], margin=margin))
+    entail_loss = torch.stack(entail, dim=0).mean()
+
+    # ---- 2. sibling separation --------------------------------------------
+    sibling = []
+    for i in range(tree.num_leaves):
+        mid_idx = tree.leaf_to_mid[i]
+        d_parent = manifold.geodesic_dist(leaves_h[:, i], mid_h[:, mid_idx])
+        d_uncle = manifold.geodesic_dist(leaves_h[:, i], mid_h[:, 1 - mid_idx])
+        sibling.append(F.relu(d_parent - d_uncle + sibling_margin))
+    sibling_loss = torch.stack(sibling, dim=0).mean()
+
+    # ---- 3. radius ordering: root < mid < leaf ------------------------------
+    r_root = manifold.radius(root)
+    r_mid = manifold.radius(mid_h)
+    r_leaf = manifold.radius(leaves_h)
+
+    radius = [F.relu(r_root.unsqueeze(-1) - r_mid + radius_margin).mean()]
+    for i in range(tree.num_leaves):
+        parent_r = r_mid[:, tree.leaf_to_mid[i]]
+        radius.append(F.relu(parent_r - r_leaf[:, i] + radius_margin).mean())
+    radius_loss = torch.stack(radius, dim=0).mean()
+
+    return w_entail * entail_loss + w_sibling * sibling_loss + w_radius * radius_loss
+
+
+def alignment_loss(
+    prompt_h: torch.Tensor,
+    visual_h: torch.Tensor,
+    manifold: LorentzManifold,
+    margin: float = 0.0,
+) -> torch.Tensor:
+    """
+    Align prompt hierarchy nodes to the batch-mean visual hierarchy.
+
+    Uses entailment rather than raw geodesic distance: minimizing distance alone
+    drives every node onto the same point, which destroys the hierarchy it is
+    supposed to preserve. Here each visual node must entail its matching prompt
+    node, which pulls them together only along the cone axis.
+    """
+    visual_mean = visual_h.mean(dim=0, keepdim=True)
+    if prompt_h.dim() == 2:
+        prompt_h = prompt_h.unsqueeze(0)
+    n = min(prompt_h.shape[1], visual_mean.shape[1])
+    prompt_nodes = prompt_h[:, :n]
+    visual_nodes = visual_mean[:, :n]
+    if prompt_nodes.shape[0] != visual_nodes.shape[0]:
+        visual_nodes = visual_nodes.expand(prompt_nodes.shape[0], -1, -1)
+    return manifold.entailment(visual_nodes, prompt_nodes, margin=margin).mean()
+
+
+def attribute_entailment_loss(
+    node_h: torch.Tensor,
+    attr_h: torch.Tensor,
+    groups: dict[int, list[int]],
+    manifold: LorentzManifold,
+    margin: float = 0.0,
+) -> torch.Tensor:
+    """
+    Ground the visual hierarchy in the attribute label semantics (HyCoCLIP-style
+    compositional entailment).
+
+    Each hierarchy node must entail the text embeddings of the attributes that
+    belong to it, e.g. the root entails the whole-person attributes (age,
+    gender) while the lower-body leaf entails ``lowerBody*``. This gives the
+    hyperbolic geometry a label-grounded objective instead of a purely
+    self-supervised one.
+
+    Args:
+        node_h: [B, N, D+1] hierarchy nodes in Lorentz coords.
+        attr_h: [A, D+1] attribute text embeddings in Lorentz coords.
+        groups: node index -> list of attribute indices.
+    """
+    if not groups:
+        return node_h.new_zeros(())
+
+    losses = []
+    b = node_h.shape[0]
+    for node_idx, attr_ids in groups.items():
+        if not attr_ids or node_idx >= node_h.shape[1]:
+            continue
+        parent = node_h[:, node_idx].unsqueeze(1).expand(-1, len(attr_ids), -1)
+        child = attr_h[attr_ids].unsqueeze(0).expand(b, -1, -1)
+        losses.append(manifold.entailment(parent, child, margin=margin).mean())
+
+    if not losses:
+        return node_h.new_zeros(())
+    return torch.stack(losses, dim=0).mean()
+
+
+# Backwards-compatible alias; the margin-only formulation it used to implement
+# could not express a hierarchy (no radius ordering, no cones).
+hierarchical_contrastive_loss = hierarchical_entailment_loss

requirements.txt

@@ -0,0 +1 @@
+torch>=2.0.0