| --- |
| tags: |
| - sentence-transformers |
| - sentence-similarity |
| - feature-extraction |
| - mathematics |
| - scientific-papers |
| - retrieval |
| - matryoshka |
| base_model: allenai/specter2_base |
| library_name: sentence-transformers |
| pipeline_tag: sentence-similarity |
| language: en |
| license: apache-2.0 |
| --- |
| |
| # math-embed |
|
|
| A 768-dimensional embedding model fine-tuned for mathematical document retrieval, with a focus on **combinatorics** and related areas (representation theory, symmetric functions, algebraic combinatorics). Built on [SPECTER2](https://huggingface.co/allenai/specter2_base) and trained using knowledge-graph-guided contrastive learning. |
|
|
| ## Performance |
|
|
| Benchmarked on mathematical paper retrieval (108 queries, 4,794 paper chunks): |
|
|
| | Model | MRR | NDCG@10 | |
| |-------|-----|---------| |
| | **math-embed (this model)** | **0.816** | **0.736** | |
| | OpenAI text-embedding-3-small | 0.461 | 0.324 | |
| | SPECTER2 (proximity adapter) | 0.360 | 0.225 | |
| | SciNCL | 0.306 | 0.205 | |
|
|
| ## Usage |
|
|
| ```python |
| from sentence_transformers import SentenceTransformer |
| |
| model = SentenceTransformer("RobBobin/math-embed") |
| |
| # Embed queries and documents |
| queries = ["Kostka polynomials", "representation theory of symmetric groups"] |
| docs = ["We study the combinatorial properties of Kostka numbers..."] |
| |
| query_embs = model.encode(queries) |
| doc_embs = model.encode(docs) |
| ``` |
|
|
| ### Matryoshka dimensions |
|
|
| Trained with Matryoshka Representation Learning — you can truncate embeddings to smaller dimensions (512, 256, 128) with graceful degradation: |
|
|
| ```python |
| # Use 256-dim embeddings for faster retrieval |
| embs = model.encode(texts) |
| embs_256 = embs[:, :256] |
| ``` |
|
|
| ## Training |
|
|
| ### Method |
| - **Loss**: MultipleNegativesRankingLoss + MatryoshkaLoss |
| - **Training data**: 22,609 (anchor, positive) pairs generated from a knowledge graph of mathematical concepts |
| - **Direct pairs**: concept name/description → chunks from that concept's source papers |
| - **Edge pairs**: cross-concept pairs from knowledge graph edges (e.g., "generalizes", "extends") |
| - **Base model**: `allenai/specter2_base` (SciBERT pre-trained on 6M citation triplets) |
|
|
| ### Configuration |
| - Epochs: 3 |
| - Batch size: 8 (effective 32 with gradient accumulation) |
| - Learning rate: 2e-5 |
| - Max sequence length: 256 tokens |
| - Matryoshka dims: [768, 512, 256, 128] |
|
|
| ### Model lineage |
| ``` |
| BERT (Google, 110M params) |
| └─ SciBERT (Allen AI, retrained on scientific papers) |
| └─ SPECTER2 base (Allen AI, + 6M citation triplets) |
| └─ math-embed (this model, + KG-derived concept-chunk pairs) |
| ``` |
|
|
| ## Approach |
|
|
| The knowledge graph was constructed by an LLM (GPT-4o-mini) from 75 mathematical research papers, identifying 559 concepts and 486 relationships. This graph provides structured ground truth: each concept maps to specific papers, and those papers' chunks serve as positive training examples. |
|
|
| This is a form of **knowledge distillation** — a large language model's understanding of mathematical relationships is distilled into a small, fast embedding model suitable for retrieval. |
|
|
| ## Limitations |
|
|
| - Trained specifically on combinatorics papers (symmetric functions, representation theory, partition identities, algebraic combinatorics) |
| - May not generalize well to other areas of mathematics or other scientific domains without additional fine-tuning |
| - 256-token context window (standard for BERT-based models) |
|
|
| ## Citation |
|
|
| See the accompanying paper: [*Knowledge-Graph-Guided Fine-Tuning of Embedding Models for Mathematical Document Retrieval*](https://huggingface.co/RobBobin/math-embed/blob/main/paper/math_embeddings.pdf) |
|
|
