neembed

Fine-tune pretrained sentence embedding models in non-Euclidean spaces with Geoopt.

Documentation · 日本語

Status: Package version v0.9.0 adds explicit Euclidean, SphereProjection, and fixed signed-curvature Stereographic geometry, plus a float64 stereographic stability policy and matched constant-curvature regression example, while preserving all v0.4-v0.8 public paths and the existing Poincaré/Lorentz curvature contract. The API remains intentionally small and may still evolve before a stable 1.0 release.

neembed is a lightweight integration layer between pretrained Sentence Transformer models and manifold-valued representations. It keeps the pretrained encoder intact, optionally projects its Euclidean output, and delegates constant-curvature geometry to Geoopt.

Pretrained Sentence Encoder
        ↓
Euclidean sentence embedding
        ↓
Projection head (optional)
        ↓
Tangent-space representation
        ↓
Geoopt manifold map
        ↓
Non-Euclidean embedding

Why non-Euclidean embeddings?

Hierarchical and tree-like relations can be awkward to represent in a flat Euclidean space. Hyperbolic spaces are a natural fit for rapidly expanding structures such as:

• taxonomies and concept hierarchies
• knowledge graphs
• hierarchical labels
• tree-like semantic relations

The current API supports Euclidean, Poincaré, Lorentz / Hyperboloid, SphereProjection, and generic signed-curvature Stereographic geometry through the same sentence-model, objective, evaluator, retrieval, and save/load surface. Geometry choice is task-dependent; neembed does not claim that curvature alone guarantees quality gains.

Current scope

v0.4 keeps the fixed-curvature v0.3 path backward-compatible:

• Sentence Transformers as the pretrained encoder backend
• Poincaré-ball and Lorentz / Hyperboloid embeddings through Geoopt
• optional lower-dimensional tangent-space projection
• shared public curvature semantics across the two hyperbolic models
• geodesic distance and manifold-aware multiple-negatives ranking loss
• ordinary AdamW fine-tuning for the model-only path
• encode() / distance(), ManifoldEmbeddingEvaluator, DataLoader interoperability, and local save/load
• geometry-consistency regressions and a matched Euclidean-vs-Poincaré-vs-Lorentz engineering benchmark

v0.4 also adds:

• opt-in fixed vs learnable curvature for Poincaré and Lorentz
• true trainable manifold-valued ManifoldPrototypes
• ManifoldPrototypeHierarchyLoss for sentence assignments plus parent-child structure
• an explicit caller-supplied Geoopt Riemannian optimizer path for manifold parameters
• joint learnable-curvature / prototype training through Geoopt stabilization
• a compact fixed-vs-learnable structure regression example

v0.5 adds a focused retrieval workflow without introducing a retrieval framework:

• optional caller-supplied explicit hard negatives while preserving the original (anchors, positives) training contract
• Recall@K and MRR on the aligned retrieval evaluator
• model.rank() for small in-memory geodesic reranking
• nearest-prototype assignment evaluation for learned manifold structure
• one reproducible v0.5 retrieval regression example

v0.6 extends that workflow to exact caller-owned corpora without adding ANN infrastructure:

• chunked exact geodesic distance evaluation without materializing the full query-by-corpus distance matrix
• exact_corpus_search() for exact text-corpus top-k retrieval
• ManifoldCorpusRetrievalEvaluator with explicit query/corpus IDs and multi-positive Recall@K / MRR
• deterministic caller-invoked mine_hard_negatives() with positive, self, and additional exclusions
• a reproducible v0.6 exact-search → evaluation → mining → explicit-negative training regression example

v0.7 expands the retrieval-objective and evaluation surface without changing the v0.4-v0.6 contracts:

• ManifoldTripletLoss for aligned geodesic margin triplets
• ManifoldMarginMSELoss for teacher positive-minus-negative margin regression
• ManifoldDistanceMSELoss for direct geodesic-distance regression
• opt-in ManifoldSymmetricMultipleNegativesRankingLoss while keeping the original one-directional MNRL unchanged
• ManifoldGradedCorpusRetrievalEvaluator with nDCG@K while preserving the existing binary/multi-positive evaluator output
• a deterministic Poincaré objective-comparison example covering MNRL, Triplet, MarginMSE, DistanceMSE, MRR, Recall@K, and nDCG@K

v0.8 adds explicit hierarchy-native learning without introducing a graph framework:

• caller-owned node IDs, parent-child edges, optional depths, and directed negatives
• ManifoldRadialOrderLoss, ManifoldDepthLoss, and ManifoldHierarchyTripletLoss
• ManifoldRetrievalHierarchyLoss for one retrieval objective plus one hierarchy objective
• ManifoldHierarchyEvaluator for radial-order and depth-vs-radius structure diagnostics
• a deterministic Poincaré retrieval-only vs hierarchy-aware regression example

v0.9 adds constant-curvature geometry selection without changing the legacy hyperbolic curvature meaning:

• Euclidean as the explicit flat K=0 backend
• SphereProjection for fixed positive sectional curvature
• generic Stereographic for fixed signed sectional curvature (K<0, K=0, or K>0)
• legacy Poincaré/Lorentz curvature=c>0 remains the magnitude of K=-c
• float64 geometry operations for SphereProjection/Stereographic, including CPU geometry fallback when the encoder runs on Apple MPS
• a deterministic matched constant-curvature comparison example for engineering regression, not geometry ranking

Product/mixed-curvature models, learnable signed curvature crossing zero, advanced SPD/Siegel/Stiefel/Birkhoff geometry, ontology parsing, graph-database integration, ANN/vector-database integration, and distributed retrieval remain outside this scope.

A manifold-valued output does not by itself require Riemannian optimization: encoder/projection parameters and learnable curvature are not manifold-valued points. Detailed parameter, optimizer, persistence, and numerical behavior lives in the Learnable structure guide. The end-to-end retrieval composition and the distinction between small in-memory reranking, exact corpus search, and external ANN systems are documented in the Retrieval workflow guide. Objective and graded-evaluation selection for v0.7 is documented in the Retrieval objectives guide. Explicit v0.8 hierarchy supervision, origin/radius semantics, composition, and structure metrics are documented in the Hierarchy-native learning guide. v0.9 geometry roles, curvature naming, compatibility, dtype/device behavior, and matched-comparison guidance are documented in the Constant-curvature geometry guide.

Installation

pip install neembed-geoopt

The PyPI distribution is named neembed-geoopt; the Python import package remains neembed.

For development:

git clone https://github.com/t-yamsaki/neembed.git
cd neembed
pip install -e ".[dev]"

Quick start

from neembed import (
    ManifoldSentenceTransformer,
    ManifoldMultipleNegativesRankingLoss,
    ManifoldTrainer,
)

model = ManifoldSentenceTransformer(
    "sentence-transformers/all-MiniLM-L6-v2",
    manifold="poincare",
    embedding_dim=64,
    curvature=1.0,
)

loss = ManifoldMultipleNegativesRankingLoss(
    model=model,
    temperature=0.1,
)

trainer = ManifoldTrainer(model=model, loss=loss)

train_batches = [
    (["Shiba Inu", "Siamese cat"], ["dog", "cat"]),
    (["dog", "cat"], ["mammal", "feline"]),
]

trainer.fit(train_batches, epochs=1)

embeddings = model.encode(["Shiba Inu", "dog", "mammal"])
distance = model.distance(embeddings[0], embeddings[1])
print(float(distance))

Each anchor is paired with the positive at the same batch index. Because off-diagonal candidates become in-batch negatives, avoid duplicate positives within one batch. This model-only path keeps the ordinary AdamW behavior even though its outputs lie on a manifold. See the Training guide for the objective and batching details, the Retrieval workflow guide for optional explicit negatives and retrieval evaluation, the Retrieval objectives guide for v0.7 objective/metric selection, the Hierarchy-native learning guide for v0.8 explicit hierarchy supervision, the Constant-curvature geometry guide before switching v0.9 geometry/curvature semantics, and the Learnable structure guide before adding trainable manifold prototypes.

Documentation

The full guide is hosted on Read the Docs:

• Installation
• Quick Start
• Architecture
• Constant-curvature geometry
• Learnable structure
• Training
• Retrieval workflow
• Retrieval objectives
• Hierarchy-native learning
• Evaluation
• Inference
• Saving and Loading
• API Reference

Examples and validation

Run the main references from the repository root:

python examples/train_poincare.py
python examples/train_lorentz.py
python examples/v04_learnable_structure.py
python examples/v05_retrieval_workflow.py
python examples/v06_exact_retrieval_workflow.py
python examples/v07_objective_comparison.py
python examples/v08_hierarchy_learning.py
python examples/v09_constant_curvature_comparison.py

• examples/train_poincare.py — minimal Poincaré workflow
• examples/train_lorentz.py — Lorentz train/evaluate/inference with intrinsic-vs-ambient dimensions
• examples/train_dataloader.py — ordinary PyTorch DataLoader training and epoch validation
• examples/train_hierarchy.py — minimal hierarchy-aware prototype objective
• examples/v04_learnable_structure.py — fixed-vs-learnable structure regression diagnostics; not a superiority benchmark
• examples/v05_retrieval_workflow.py — explicit hard negatives, Recall@K / MRR, in-memory geodesic reranking, and prototype assignment in one Poincaré regression workflow; not a research benchmark
• examples/v06_exact_retrieval_workflow.py — exact corpus search, explicit-ID corpus evaluation, offline hard-negative mining, and the existing three-sequence trainer in one Poincaré regression workflow; not a research benchmark
• examples/v07_objective_comparison.py — deterministic comparison of MNRL, Triplet, MarginMSE, and DistanceMSE under fixed data/initialization with MRR, Recall@K, and nDCG@K diagnostics; not a research benchmark or superiority claim
• examples/v08_hierarchy_learning.py — deterministic retrieval-only vs hierarchy-aware Poincaré regression with explicit caller-owned hierarchy supervision and separate retrieval/structure diagnostics; not a benchmark or superiority claim
• examples/v09_constant_curvature_comparison.py — deterministic matched comparison of Euclidean, Poincaré, Lorentz, SphereProjection, and signed Stereographic geometry with retrieval and dtype/device diagnostics; engineering regression only, not a benchmark or geometry-superiority claim
• experiments/README.md — reproducible Euclidean-vs-Poincaré-vs-Lorentz engineering benchmark and interpretation limits

License

neembed is released under the MIT License. Third-party dependencies retain their own licenses.