neembed

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

Documentation · 日本語

Status: Package version v0.11.0 adds reproducible encoder loading, differentiable query/document input roles, batched encoding, versioned checkpoint metadata, and documented public API/benchmark contracts. It retains v0.10.0 ProductManifold/Scaled mixed-curvature support and all v0.4-v0.9 single-manifold contracts. 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

v0.10 adds flat mixed-curvature products with fixed component scales, full-product retrieval, named distance diagnostics, and optional hierarchy supervision on a designated Poincaré/Lorentz component. See the Mixed-curvature product guide for configuration, compatibility, optimizer/persistence behavior, and the regression example. More geometry does not guarantee better task metrics.

Nested products, learnable product curvature/scales, 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.

Role-aware training and search

After installing the development checkout, download the two pinned encoders once with HF_HUB_DISABLE_XET=1 python scripts/preload_role_models.py. The following example then runs from the repository root with cached files only. E5-small-v2 uses English text with explicit query/passage instructions; pass raw texts.

import torch
from neembed import ManifoldMultipleNegativesRankingLoss, ManifoldSentenceTransformer, exact_corpus_search

model = ManifoldSentenceTransformer(
    "intfloat/e5-small-v2", revision="ffb93f3bd4047442299a41ebb6fa998a38507c52",
    local_files_only=True, device="cpu", embedding_dim=16,
)
model.encoder.prompts.update({"query": "query: ", "document": "passage: "})
queries = ["What is a dog?", "Which animals can fly?"]
documents = ["A dog is a mammal.", "Birds can fly.", "Trees are plants."]

loss_fn = ManifoldMultipleNegativesRankingLoss(model, input_options={
    "anchors": {"task": "query"}, "positives": {"task": "document"},
})
model.train()
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
optimizer.zero_grad()
loss = loss_fn(queries, documents[:2])
loss.backward()
optimizer.step()

embeddings = model.encode_query(queries, batch_size=1, convert_to_tensor=True)
results = exact_corpus_search(model, queries, documents, top_k=2, input_options={
    "queries": {"task": "query"}, "corpus": {"task": "document"},
})
model.save_pretrained("./saved_role_model")
restored = ManifoldSentenceTransformer.from_pretrained("./saved_role_model", local_files_only=True, device="cpu")
torch.testing.assert_close(restored.encode_query(queries, batch_size=1, convert_to_tensor=True), embeddings)
print(results)

forward and losses retain the caller's gradient context and model mode. encode switches to eval mode and returns CPU NumPy arrays by default, or geometry-device tensors with convert_to_tensor=True; call model.train() when resuming training. distance is a no-grad helper with broadcast Tensor output and leaves mode unchanged. Use model.manifold.dist on forward outputs for differentiable distances. Rankings put the smallest geodesic distance first; equal distances keep candidate/corpus order.

For intrinsic dimension D, single geometry outputs have width D, except Lorentz's D + 1. Product width is product_config.ambient_dim, the sum of component intrinsic dimensions plus one for each Lorentz component. Lorentz, SphereProjection, Stereographic, and fixed-double products map/output float64; Poincare/Euclidean retain ordinary dtype. Tensors follow the geometry device (SphereProjection/Stereographic and fixed-double products fall back to CPU on MPS); NumPy outputs use CPU. NumPy conversion must support the output dtype: bfloat16 requires Tensor output. These policies do not qualify every device/dtype.

Unspecified roles keep raw-text behavior even with a saved default prompt. Versionless v0.10-and-earlier checkpoints still load; new format-version 1 metadata preserves prompts, geometry, dtype and available provenance. Per-call options and batch size, optimizer state, and external prototypes are not saved as defaults. Float64 encoders include supplemental exact weights. See the public API contracts and compatibility guide.

Documentation

Local saves now use versioned metadata for geometry, prompts, floating dtypes, package version, and available base-model provenance. Versionless v0.10 and earlier configurations remain loadable; restoration uses the saved local encoder.

The minimum qualified dependencies are PyTorch 2.3.0, Sentence Transformers 5.4.1, and Geoopt 0.5.1. A two-profile CPU CI verifies pinned MiniLM and E5-small-v2 encoders, including role-aware training, retrieval, and offline save/load. See the compatibility guide for exact stacks, model revisions, and tested scope.

encode(), encode_query(), and encode_document() accept batch_size=32 by default, preserving input order and geometry dtype/device. Set a smaller positive integer to bound each encoder call; empty sequences return an empty matrix with the configured output width. Encoding leaves the model in eval mode.

forward() and encode() accept task="query" / "document", prompt_name, or an explicit prompt. encode_query() and encode_document() share the same prompt selection with the differentiable training path. Calls without these options retain raw-text behavior. See the Inference guide for precedence, prompt-aware pooling, routing support, and current integration limits.

Built-in losses accept input_options, mapping their text input names to the same task/prompt options. Configure anchors as queries and positives/explicit negatives as documents while keeping existing training batches. See the Retrieval objectives guide for examples and independent hierarchy/composite input settings.

rank(), exact corpus search, evaluators, and offline hard-negative mining also accept input_options. Reuse the same query/document settings as training, using each API's text input names. See the Retrieval workflow guide for a complete example and the input-name mapping.

Encoder loading accepts revision, local_files_only, cache_folder, and device in both the constructor and local from_pretrained(). Pin a remote encoder with a commit ID, or use local_files_only=True for local/cached files. See the Saving and Loading guide for examples and the distinction between encoder revision and local checkpoints.

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
python examples/v10_mixed_curvature_workflow.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
• examples/v10_mixed_curvature_workflow.py — fixed-scale Poincaré × SphereProjection × Euclidean retrieval plus named hyperbolic-component hierarchy supervision, component distances, and save/load validation; not a benchmark or mixed-curvature 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.