# Matryoshka Text Embedder

> Kusupati et al. 2022 - one embedding model whose prefixes are independently usable, trained with InfoNCE on every prefix at once.

Pick when callers have different index budgets: truncate to 96 dimensions instead of training and serving a second, smaller model. Adds no parameters and no inference cost.

- Category: Multimodal
- Paper: Kusupati et al. (2022). Matryoshka Representation Learning. NeurIPS 2022 (https://arxiv.org/abs/2205.13147)
- Layers: 11
- Parameters: 54.74M
- Input shape (batchless): 1 × 512
- Output shape: 1 × 768
- Verifier verdict: pass
- Graph JSON: https://neurarch.com/templates/matryoshka-embed/model.json
- Open on the canvas: https://neurarch.com/?template=matryoshka-embed

## Structure

| # | Layer | Type | Parameters | Output shape |
| --- | --- | --- | --- | --- |
| 1 | Text tokens | Input | shape=[1, 512] | 1 × 512 |
| 2 | Token Embed | Embedding | vocabSize=30522 | 1 × 512 × 768 |
| 3 | Pos Encoding | Positional Encoding | maxLen=512 | 1 × 512 × 768 |
| 4 | Encoder Block 1 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 1 × 512 × 768 |
| 5 | Encoder Block 2 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 1 × 512 × 768 |
| 6 | Encoder Block 3 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 1 × 512 × 768 |
| 7 | Encoder Block 4 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 1 × 512 × 768 |
| 8 | Pooling | Attention Pooling (PMA) | numHeads=12 | 1 × 768 |
| 9 | Contrastive Head (InfoNCE) | Contrastive Head (InfoNCE / CLIP) |  | 1 × 768 |
| 10 | Matryoshka (768/384/192/96/48) | Matryoshka Head (MRL) | embedDim=768 | 1 × 768 |
| 11 | embedding (truncatable) | Output |  | 1 × 768 |

## Verifier findings

- **info** `deep-no-norm`: 9 layers with no BatchNorm, LayerNorm, or GroupNorm. Without normalization, activations can explode or vanish across layers, causing slow or unstable training. Fix: Add BatchNorm after Conv2d (CV tasks), LayerNorm after attention/FFN (NLP/LLM), or GroupNorm for small batch sizes.

## Exported PyTorch (first 46 lines)

```python
# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class AttentionPooling(nn.Module):
    """Pooling by Multihead Attention (Set Transformer's PMA). Learned seed
    vectors attend over the sequence, so the sequence axis is consumed and the
    pooling is learned rather than an unweighted mean."""

    def __init__(self, dim: int = 512, num_heads: int = 8, num_seeds: int = 1):
        super().__init__()
        self.seeds = nn.Parameter(torch.randn(num_seeds, dim) * 0.02)
        heads = num_heads if dim % num_heads == 0 else 1
        self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.num_seeds = num_seeds

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        q = self.seeds.unsqueeze(0).expand(x.size(0), -1, -1)
        out = self.attn(q, x, x, need_weights=False)[0]
        return out.squeeze(-2) if self.num_seeds == 1 else out


class ContrastiveHead(nn.Module):
    """CLIP-style projection into the shared space: bias-free, L2-normalised,
    with the temperature as a learned parameter. After the normalisation a dot
    product is a cosine and is bounded, which is what makes the temperature
    mean anything."""

    def __init__(self, d_in: int, proj_dim: int = 512,
                 temperature: float = 0.07, learnable_temp: bool = True,
                 normalize: bool = True):
        super().__init__()
        self.proj = nn.Linear(d_in, proj_dim, bias=False)
        self.normalize = normalize
        scale = torch.tensor(float(1.0 / max(temperature, 1e-6))).log()
        self.logit_scale = nn.Parameter(scale) if learnable_temp else scale

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        z = self.proj(x)
        return F.normalize(z, dim=-1) if self.normalize else z
```

## Machine access

- Every architecture: https://neurarch.com/a/index.json
- Verify a graph of your own: `POST https://www.neurarch.com/api/v1/check` (see https://neurarch.com/developer.html)
- MCP server, so an agent edits the graph with the checks in the loop: https://neurarch.com/docs/mcp.md