Architectures / Multimodal

🪆 Matryoshka Text Embedder

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

From Kusupati et al. (2022). Matryoshka Representation Learning. NeurIPS 2022. This page is the graph, not the PDF: open it, edit it, verify it, export it.

Every number on this page is computed from the graph by the same functions the app runs, not written by hand.

When to pick it

Structure

11 layers. Output shapes are propagated from the input shape, batch dimension excluded.

What the verifier says

The same 43 structural checks that run on every edit in the app, on this graph.

The PyTorch it exports

Generated from the graph above. First 46 lines; the app exports the whole file, plus the training loop, the data contract and a deploy bundle.

# 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

For agents

Also in Multimodal