README.md

---
language:
  - zh
  - en
license: apache-2.0
library_name: sentence-transformers
tags:
  - text-embedding
  - sentence-transformers
  - feature-extraction
  - chinese
  - minimind
  - qwen3-embedding
  - matryoshka
  - moe
base_model: jingyaogong/minimind-3
pipeline_tag: feature-extraction
---

<div align="center">

<img src="./images/banner.png" alt="MiniMind-Embedding" width="100%">

</div>


# MiniMind-Embedding-MoE

<div align="center">

![params](https://img.shields.io/badge/params-198M_A64M-blue)
![dim](https://img.shields.io/badge/dim-768-green)
![max_seq](https://img.shields.io/badge/max__seq-8192-orange)
![STS](https://img.shields.io/badge/C_MTEB%20STS-0.42-yellow)

</div>

A Mixture-of-Experts (MoE) text embedding model built on the [MiniMind-3-MoE](https://huggingface.co/jingyaogong/minimind-3-pytorch) backbone (198M total / 64M active params), with the training pipeline fully aligned with [Qwen3-Embedding](https://arxiv.org/abs/2506.05176). This is the **MoE counterpart** of [MiniMind-Embedding-Dense](https://huggingface.co/Muzian/minimind-embedding-dense).

## 📌 Introduction

This model uses the same Qwen3-Embedding recipe as the Dense version (causal attention + last-token pooling + InfoNCE with false-negative masking + MRL), but swaps the feed-forward layers for a **4-expert top-1 MoE**. The goal is to study whether MoE architecture helps embedding quality at the small scale.

### Core techniques

<div align="center">

![Last-Token Pooling](./images/concept_last_token_pool.png)

</div>

<div align="center">

![InfoNCE + False-Negative Mask](./images/concept_false_neg_mask.png)

</div>

<div align="center">

![Matryoshka MRL](./images/concept_mrl.png)

</div>

## 🏗️ Model Details

| Item | Value |
|------|-------|
| Architecture | MiniMind-3-MoE (causal decoder-only, 4 experts top-1) |
| Total params | 198M |
| Active params | 64M (per token, 1 expert active) |
| Hidden size | 768 |
| Layers | 8 |
| Attention heads | 8 (query) / 4 (KV, GQA) |
| Vocab size | 6400 |
| Embedding output dim | 768 (MRL: truncatable to 512/256/128/64) |
| Max sequence length | 8192 |
| Pooling | last-token |

## 🎓 Training Details (3 Stages)

| Stage | Data | Loss | Steps | Final Loss | Status |
|-------|------|------|-------|------------|--------|
| **1. Weakly-supervised** | T2Ranking triplets (90k) | Standard InfoNCE | 2,827 | ~0.45 | ✅ Done |
| **2. Supervised** | T2Ranking-15 (340k, 15 hard negs) | InfoNCE + false-neg mask + MRL | 63,768 (3 epochs) | ~0.85 | ✅ Done |
| **3. Model merging** | Last 5 Stage-2 checkpoints | SLERP | — | — | ✅ Done |

<div align="center">

![MoE Training Loss](./images/moe_training_loss.png)

*MoE Stage 1 (left) + Stage 2 (right), 3 epochs.*

</div>

### Dense vs MoE Training Loss

<div align="center">

![Dense vs MoE Loss](./images/dense_vs_moe_loss.png)

</div>

> Note: MoE loss is consistently higher than Dense. This is expected — MoE routes tokens to different experts, making the InfoNCE contrastive task harder (each expert sees only a subset of the sequence's tokens). The higher train loss doesn't directly translate to worse embeddings, but the evaluation shows MoE underperforms Dense at this scale.

## 📊 Evaluation

### C-MTEB STS (Spearman correlation)

<div align="center">

![Dense vs MoE Comparison](./images/dense_vs_moe.png)

</div>

| Task | Dense (64M) | **MoE (198M-A64M)** | Δ |
|------|:-----------:|:-------------------:|:---:|
| ATEC | 0.265 | **0.201** | -0.064 |
| BQ | 0.379 | **0.297** | -0.082 |
| LCQMC | 0.631 | **0.582** | -0.049 |
| STSB | 0.637 | **0.607** | -0.030 |
| **Average** | **0.478** | **0.422** | **-0.056** |

### Key Finding: MoE underperforms Dense at small scale

> 💡 **Surprising result**: the MoE model (198M total / 64M active) **underperforms** the Dense model (64M) by 0.056 average STS. This is a counterintuitive but important finding:
>
> 1. **MoE fragmentation**: with only 4 experts and top-1 routing, each expert sees only ~25% of tokens. For embedding (where the whole sentence's representation matters), this fragmentation hurts — no single expert gets a complete view of the sentence.
> 2. **Active params are the same**: MoE activates 64M params per token (same as Dense's 64M), so it's not "bigger" in terms of compute — it's just more fragmented.
> 3. **Routing overhead**: the router must learn to dispatch tokens, adding an optimization burden that competes with learning good representations.
>
> **Takeaway**: for small embedding models, Dense > MoE. MoE benefits emerge at larger scales (Qwen3-Embedding uses MoE only at 4B+ params) where each expert is large enough to be expressive on its own.

## 🚀 Usage

```python
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

model = SentenceTransformer("Muzian/minimind-embedding-moe")

documents = ["天空呈蓝色是因为阳光散射。", "红烧肉的做法是..."]
doc_embs = model.encode(documents, normalize_embeddings=True)

queries = ["天空为什么是蓝色的"]
query_embs = model.encode(queries, prompt_name="query", normalize_embeddings=True)

scores = cos_sim(query_embs, doc_embs)
```

> **Recommendation**: for most use cases, prefer [MiniMind-Embedding-Dense](https://huggingface.co/Muzian/minimind-embedding-dense) — it's smaller, faster, and scores higher. This MoE model is primarily for research/comparison.

## ⚠️ Limitations

1. Same backbone limitations as Dense (vocab=6400, limited Stage-1 data)
2. **MoE underperforms Dense** at this scale — see analysis above
3. 198M params but only 64M active → inference is slower than Dense without quality gain

## 🙏 Acknowledgements

- Base: [MiniMind](https://github.com/jingyaogong/minimind) by @jingyaogong
- Route: [Qwen3-Embedding](https://arxiv.org/abs/2506.05176)
- Data: [T2Ranking](https://huggingface.co/datasets/THUIR/T2Ranking)
- Eval: [MTEB](https://github.com/embeddings-benchmark/mteb)

## 📄 License

Apache 2.0