Efficient Mixture of Experts Text Embeddings

# Illustrative sketch (not from the paper)
import torch
import torch.nn as nn
import torch.nn.functional as F

# 1. Fully visible bidirectional mask for decoder‑only model
def bidirectional_mask(seq_len):
    # True where attention is allowed (all positions)
    return torch.ones(seq_len, seq_len, dtype=torch.bool)

# 2. Sparse Mixture‑of‑Experts layer with top‑4 routing (placeholder)
class MoE(nn.Module):
    def __init__(self, expert_dim, num_experts, top_k=4):
        super().__init__()
        self.experts = nn.ModuleList([nn.Linear(expert_dim, expert_dim) for _ in range(num_experts)])
        self.top_k = top_k
        self.gate = nn.Linear(expert_dim, num_experts)
    def forward(self, x):
        # x: [batch, seq, dim]
        gate_scores = self.gate(x)                     # [batch, seq, num_experts]
        topk_vals, topk_idx = torch.topk(gate_scores, self.top_k, dim=-1)
        # Simple routing: sum selected expert outputs
        out = torch.zeros_like(x)
        for i in range(self.top_k):
            idx = topk_idx[..., i]
            expert_out = torch.stack([self.experts[e](x[b]) for b, e in enumerate(idx)])
            out += expert_out
        return out

# 3. Encoder that applies mask, MoE, mean pooling, and L2 normalisation
class GigaEncoder(nn.Module):
    def __init__(self, dim, num_experts):
        super().__init__()
        self.moe = MoE(dim, num_experts)
    def forward(self, token_states, mask):
        # token_states: [batch, seq, dim]
        h = self.moe(token_states)                     # sparse expert processing
        # mean pooling over visible tokens
        pooled = (h * mask.unsqueeze(-1).float()).sum(dim=1) / mask.sum(dim=1, keepdim=True).float()
        # L2 normalisation for cosine similarity
        return F.normalize(pooled, p=2, dim=-1)

# 4. Contrastive loss (placeholder for three‑stage training)
def contrastive_loss(q, k, temperature=0.05):
    # q, k: [batch, dim] already normalised
    logits = torch.matmul(q, k.t()) / temperature
    labels = torch.arange(q.size(0), device=q.device)
    return F.cross_entropy(logits, labels)

# 5. Similarity‑distribution distillation loss
def distillation_loss(student_emb, teacher_emb):
    # Compute pairwise cosine similarity matrices
    s_sim = torch.matmul(student_emb, student_emb.t())
    t_sim = torch.matmul(teacher_emb, teacher_emb.t())
    # Convert to probability distributions (softmax over rows)
    s_prob = F.log_softmax(s_sim, dim=-1)
    t_prob = F.softmax(t_sim, dim=-1)
    # KL divergence summed over rows
    return F.kl_div(s_prob, t_prob, reduction='batchmean')

// Illustrative sketch (not from the paper)
const tf = require('@tensorflow/tfjs-node');

// 1. Fully visible bidirectional mask (boolean tensor)
function bidirectionalMask(seqLen) {
  return tf.ones([seqLen, seqLen], 'bool');
}

// 2. Simple top‑4 MoE layer (placeholder implementation)
class MoE extends tf.layers.Layer {
  constructor(expertDim, numExperts, topK = 4) {
    super({});
    this.expertDim = expertDim;
    this.numExperts = numExperts;
    this.topK = topK;
    this.experts = [];
    for (let i = 0; i < numExperts; i++) {
      this.experts.push(tf.layers.dense({ units: expertDim }));
    }
    this.gate = tf.layers.dense({ units: numExperts });
  }
  call(inputs) {
    // inputs: [batch, seq, dim]
    const gateScores = this.gate.apply(inputs); // [batch, seq, numExperts]
    const { values, indices } = tf.topk(gateScores, this.topK);
    // Route to selected experts and sum their outputs
    let out = tf.zerosLike(inputs);
    for (let k = 0; k < this.topK; k++) {
      const idx = tf.squeeze(tf.slice(indices, [0, 0, k], [-1, -1, 1]), -1);
      const expertOutputs = tf.tidy(() => {
        // Gather expert for each batch element
        const batchOutputs = [];
        const batchSize = inputs.shape[0];
        for (let b = 0; b < batchSize; b++) {
          const expertIdx = idx.arraySync()[b];
          const expert = this.experts[expertIdx];
          batchOutputs.push(expert.apply(tf.slice(inputs, [b, 0, 0], [1, -1, -1])));
        }
        return tf.concat(batchOutputs, 0);
      });
      out = tf.add(out, expertOutputs);
    }
    return out;
  }
}

// 3. Encoder: MoE → mean pooling → L2 normalisation
class GigaEncoder {
  constructor(dim, numExperts) {
    this.moe = new MoE(dim, numExperts);
  }
  call(tokenStates, mask) {
    // tokenStates: [batch, seq, dim]
    const h = this.moe.call(tokenStates);
    const maskFloat = tf.cast(mask, 'float32').expandDims(-1);
    const summed = tf.sum(tf.mul(h, maskFloat), 1);
    const divisor = tf.sum(maskFloat, 1);
    const pooled = tf.div(summed, divisor);
    return tf.linalg.l2Normalize(pooled, -1);
  }
}

// 4. Contrastive loss (placeholder for three‑stage training)
function contrastiveLoss(q, k, temperature = 0.05) {
  const logits = tf.div(tf.matMul(q, k, false, true), temperature);
  const labels = tf.range(0, q.shape[0], 1, 'int32');
  return tf.losses.softmaxCrossEntropy(labels, logits);
}

// 5. Similarity‑distribution distillation loss
function distillationLoss(studentEmb, teacherEmb) {
  const sSim = tf.matMul(studentEmb, studentEmb, false, true);
  const tSim = tf.matMul(teacherEmb, teacherEmb, false, true);
  const sProb = tf.logSoftmax(sSim, -1);
  const tProb = tf.softmax(tSim, -1);
  const kl = tf.losses.kullbackLeiblerDiv(tProb, sProb);
  return tf.mean(kl);
}