02. Embedding Model Fine-Tuning — Contrastive Adapter for Titan Embeddings V2

graph TB
    subgraph "Frozen: Amazon Titan Embeddings V2 (Bedrock)"
        A["User Query: 'dark isekai manga like Berserk'"] --> B["Titan Embedding V2<br>(Managed API Call)<br>~300M params, frozen"]
        B --> C["Raw Embedding e ∈ ℝ¹⁰²⁴<br>General-purpose, not manga-tuned"]
    end

    subgraph "Trainable: Contrastive Adapter (~1M params)"
        C --> D["Linear W₁ (1024 → 512)<br>524,800 params<br>Projects to bottleneck"]
        D --> E["GELU Activation<br>Smooth non-linearity<br>Preserves negative dims"]
        E --> F["Linear W₂ (512 → 1024)<br>525,312 params<br>Projects back to original dim"]
        F --> G["LayerNorm (1024)<br>2,048 params<br>Stabilizes output magnitude"]
        G --> H["Residual Connection<br>z = LayerNorm(adapter(e)) + e"]
    end

    subgraph "Output"
        H --> I["Adapted Embedding z ∈ ℝ¹⁰²⁴<br>Manga-domain aware<br>cos(q, d) is now genre-sensitive"]
    end

    style B fill:#bbdefb
    style D fill:#fff9c4
    style E fill:#fff9c4
    style F fill:#fff9c4
    style G fill:#fff9c4

sequenceDiagram
    participant Batch as Training Batch<br>(128 triplets)
    participant Titan as Titan Embeddings V2<br>(Frozen, Bedrock API)
    participant Adapter as Contrastive Adapter<br>(1M trainable params)
    participant Loss as InfoNCE Loss<br>Calculator
    participant ANN as Hard Negative<br>Miner (ANN Index)
    participant Optim as Adam Optimizer

    Note over Batch: Epoch start: refresh hard negatives
    Batch->>ANN: All queries from training set
    ANN->>Batch: Top-K hard negatives per query

    Batch->>Titan: 128 queries + 128 positives + 128 negatives (384 texts)
    Titan->>Adapter: 384 frozen embeddings ∈ ℝ¹⁰²⁴

    Note over Adapter: Forward pass through adapter<br>e → W₁ → GELU → W₂ → LayerNorm → z<br>3 matrix multiplications per embedding

    Adapter->>Loss: 384 adapted embeddings
    Loss->>Loss: Compute pairwise cosine similarities<br>128 × 128 similarity matrix
    Loss->>Loss: Apply temperature τ=0.07<br>InfoNCE = -log(exp(s⁺/τ) / Σexp(s/τ))

    Note over Loss: Backward pass
    Loss->>Adapter: ∂L/∂z for all 384 embeddings
    Note over Adapter: Gradients flow through:<br>LayerNorm → W₂ → GELU' → W₁<br>Titan params receive NO gradient (frozen)

    Adapter->>Optim: Update 1M adapter params
    Optim->>Adapter: Adam step: θ ← θ - lr·m̂/(√v̂ + ε)

graph LR
    subgraph "Before Adapter (Titan V2 Raw)"
        A1["'dark fantasy manga'<br>(0.72, 0.31, ...)"]
        B1["'action adventure manga'<br>(0.70, 0.35, ...)"]
        C1["'comedy isekai'<br>(0.68, 0.29, ...)"]
        D1["'Berserk vol 1'<br>(0.71, 0.33, ...)"]
        E1["'One Piece vol 1'<br>(0.69, 0.34, ...)"]

        A1 -.->|"cos=0.94<br>too close!"| B1
        A1 -.->|"cos=0.91"| C1
        A1 -.->|"cos=0.93"| D1
        B1 -.->|"cos=0.95"| E1
    end

    subgraph "After Adapter (Manga-Tuned)"
        A2["'dark fantasy manga'<br>moved to dark cluster"]
        B2["'action adventure manga'<br>moved to action cluster"]
        C2["'comedy isekai'<br>moved to comedy cluster"]
        D2["'Berserk vol 1'<br>near dark fantasy"]
        E2["'One Piece vol 1'<br>near action adventure"]

        A2 -.->|"cos=0.92<br>correct!"| D2
        A2 -.->|"cos=0.71<br>separated"| B2
        A2 -.->|"cos=0.58<br>far apart"| C2
        B2 -.->|"cos=0.89"| E2
    end

graph BT
    subgraph "Loss Layer"
        L["InfoNCE Loss<br>L = -log(softmax(sim/τ))"]
    end

    subgraph "LayerNorm + Residual"
        LN["LayerNorm<br>∂L/∂z: scales ≈ 1.0<br>Residual adds identity gradient"]
    end

    subgraph "W₂ Linear (512 → 1024)"
        W2["W₂ Gradient:<br>∂L/∂W₂ = ∂L/∂z · h^T<br>Updates: 525K params<br>Magnitude: 0.8 × reference"]
    end

    subgraph "GELU Activation"
        G["GELU'(x) = Φ(x) + x·φ(x)<br>Gradient: 0.5-1.0 for most inputs<br>Near-linear for positive inputs<br>Smooth attenuation for negative"]
    end

    subgraph "W₁ Linear (1024 → 512)"
        W1["W₁ Gradient:<br>∂L/∂W₁ = (W₂^T · ∂L/∂z · GELU'(·)) · e^T<br>Updates: 524K params<br>Magnitude: 0.6 × reference"]
    end

    subgraph "Frozen Titan Embeddings"
        TT["Titan V2 Output e<br>NO gradient flows here<br>Bedrock API, no backprop"]
    end

    L --> LN
    LN --> W2
    W2 --> G
    G --> W1
    W1 --> TT

    style TT fill:#bbdefb
    style W1 fill:#fff9c4
    style G fill:#fff9c4
    style W2 fill:#fff9c4

graph TD
    subgraph "τ = 0.01 (Too Sharp)"
        A["sim(q, d⁺) = 0.85<br>sim(q, d₁⁻) = 0.82<br>sim(q, d₂⁻) = 0.60"]
        A --> B["softmax([85, 82, 60])<br>= [0.953, 0.045, 0.002]"]
        B --> C["Only d₁⁻ gets gradient<br>d₂⁻ effectively ignored<br>Misses useful signal"]
    end

    subgraph "τ = 0.07 (Our Setting)"
        D["sim(q, d⁺) = 0.85<br>sim(q, d₁⁻) = 0.82<br>sim(q, d₂⁻) = 0.60"]
        D --> E["softmax([12.1, 11.7, 8.6])<br>= [0.598, 0.366, 0.036]"]
        E --> F["Both hard negatives<br>get meaningful gradient<br>Balanced learning"]
    end

    subgraph "τ = 0.5 (Too Soft)"
        G["sim(q, d⁺) = 0.85<br>sim(q, d₁⁻) = 0.82<br>sim(q, d₂⁻) = 0.60"]
        G --> H["softmax([1.7, 1.64, 1.2])<br>= [0.381, 0.359, 0.260]"]
        H --> I["Easy negative gets<br>almost same weight as hard<br>Noisy, slow convergence"]
    end

graph LR
    subgraph "Epoch N"
        A["Train adapter<br>on current pairs"] --> B["Adapter weights<br>updated"]
        B --> C["Re-embed all<br>catalog items (50K)"]
        C --> D["Rebuild FAISS<br>ANN index"]
        D --> E["Mine new hard<br>negatives for<br>each query"]
    end

    subgraph "Epoch N+1"
        E --> F["Train with<br>refreshed negatives"]
        F --> G["Repeat..."]
    end

    subgraph "Why Refresh?"
        H["At epoch 1:<br>'comedy isekai' is hard<br>negative for 'dark isekai'"]
        I["At epoch 3:<br>adapter separates them<br>'comedy isekai' is now easy"]
        J["New hard negative:<br>'horror isekai' moves<br>closer, becomes the<br>new challenge"]
        H --> I --> J
    end

import boto3
import json
import numpy as np
from typing import List


class TitanEmbeddingClient:
    """
    Client for Amazon Titan Embeddings V2 via Bedrock.
    Returns 1024-dim embeddings. We call this for both training
    (to get frozen embeddings as adapter input) and inference
    (before passing through the adapter).
    """

    def __init__(self, region: str = "us-east-1"):
        self.client = boto3.client("bedrock-runtime", region_name=region)
        self.model_id = "amazon.titan-embed-text-v2:0"

    def embed_texts(self, texts: List[str], batch_size: int = 25) -> np.ndarray:
        """
        Embed a list of texts. Titan V2 supports single-text calls,
        so we batch manually.
        """
        all_embeddings = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            batch_embeddings = []
            for text in batch:
                response = self.client.invoke_model(
                    modelId=self.model_id,
                    body=json.dumps({
                        "inputText": text,
                        "dimensions": 1024,
                        "normalize": True,
                    }),
                )
                result = json.loads(response["body"].read())
                batch_embeddings.append(result["embedding"])
            all_embeddings.extend(batch_embeddings)

        return np.array(all_embeddings, dtype=np.float32)


    def embed_single(self, text: str) -> np.ndarray:
        """Single text embedding for inference path."""
        response = self.client.invoke_model(
            modelId=self.model_id,
            body=json.dumps({
                "inputText": text,
                "dimensions": 1024,
                "normalize": True,
            }),
        )
        result = json.loads(response["body"].read())
        return np.array(result["embedding"], dtype=np.float32)

import torch
import torch.nn as nn
import torch.nn.functional as F


class ContrastiveAdapter(nn.Module):
    """
    Lightweight adapter that transforms Titan V2 embeddings
    to be manga-domain aware.

    Architecture: e (1024) → W₁ (512) → GELU → W₂ (1024) → LayerNorm + Residual
    Parameters: ~1M (0.4% of Titan V2)
    """

    def __init__(self, embed_dim: int = 1024, bottleneck_dim: int = 512):
        super().__init__()
        self.linear1 = nn.Linear(embed_dim, bottleneck_dim)
        self.linear2 = nn.Linear(bottleneck_dim, embed_dim)
        self.layer_norm = nn.LayerNorm(embed_dim)

        # Initialize with small weights so adapter starts near identity
        nn.init.normal_(self.linear1.weight, std=0.02)
        nn.init.normal_(self.linear2.weight, std=0.02)
        nn.init.zeros_(self.linear1.bias)
        nn.init.zeros_(self.linear2.bias)

    def forward(self, embeddings: torch.Tensor) -> torch.Tensor:
        """
        Args:
            embeddings: (batch_size, 1024) — frozen Titan V2 outputs
        Returns:
            (batch_size, 1024) — adapted embeddings
        """
        # Bottleneck projection
        h = F.gelu(self.linear1(embeddings))    # (batch, 512)
        adapted = self.linear2(h)                # (batch, 1024)

        # Residual connection + LayerNorm
        output = self.layer_norm(adapted + embeddings)  # (batch, 1024)

        # L2 normalize for cosine similarity
        output = F.normalize(output, p=2, dim=-1)

        return output

class InfoNCELoss(nn.Module):
    """
    InfoNCE contrastive loss with in-batch negatives.

    For a batch of N (query, positive) pairs:
    - Each query has 1 positive (its paired document)
    - Each query has N-1 negatives (all other documents in the batch)
    - Total comparisons per query: N
    - This is equivalent to N-way classification

    With batch_size=128, each query gets 127 free negatives.
    """

    def __init__(self, temperature: float = 0.07):
        super().__init__()
        self.temperature = temperature

    def forward(
        self,
        query_embeddings: torch.Tensor,
        doc_embeddings: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            query_embeddings: (N, dim) — adapted query embeddings
            doc_embeddings: (N, dim) — adapted document embeddings
            Row i of queries matches row i of docs (positive pair)

        Returns:
            Scalar loss
        """
        # Compute all pairwise cosine similarities
        # (N, dim) × (dim, N) → (N, N)
        similarity_matrix = torch.mm(
            query_embeddings, doc_embeddings.t()
        ) / self.temperature

        # Labels: diagonal entries are positives
        # query[i] should match doc[i], so label for row i = i
        labels = torch.arange(
            similarity_matrix.size(0), device=similarity_matrix.device
        )

        # Cross-entropy on similarity matrix = InfoNCE
        loss = F.cross_entropy(similarity_matrix, labels)

        return loss

import faiss
import numpy as np


class HardNegativeMiner:
    """
    Mines hard negatives using FAISS approximate nearest neighbor search.
    Refreshed every epoch as the adapter changes the embedding space.
    """

    def __init__(self, embed_dim: int = 1024, top_k: int = 20):
        self.embed_dim = embed_dim
        self.top_k = top_k
        self.index = None

    def build_index(self, doc_embeddings: np.ndarray):
        """Build FAISS index from adapted document embeddings."""
        # Use IVF index for 50K+ documents
        n_docs = doc_embeddings.shape[0]
        n_lists = min(int(np.sqrt(n_docs)), 256)

        quantizer = faiss.IndexFlatIP(self.embed_dim)  # Inner product (cosine after L2 norm)
        self.index = faiss.IndexIVFFlat(
            quantizer, self.embed_dim, n_lists, faiss.METRIC_INNER_PRODUCT
        )
        self.index.train(doc_embeddings)
        self.index.add(doc_embeddings)
        self.index.nprobe = 10  # Search 10 nearest clusters

    def mine(
        self,
        query_embeddings: np.ndarray,
        positive_ids: np.ndarray,
        n_hard_negatives: int = 5,
    ) -> np.ndarray:
        """
        For each query, find top-K nearest docs and filter out the positive.
        Return indices of hard negatives.
        """
        # Search for top-K nearest documents
        similarities, indices = self.index.search(query_embeddings, self.top_k)

        hard_negatives = []
        for i in range(len(query_embeddings)):
            # Filter out the known positive
            neg_indices = [
                idx for idx in indices[i]
                if idx != positive_ids[i] and idx >= 0
            ]
            # Take top-N hard negatives
            hard_negatives.append(neg_indices[:n_hard_negatives])

        return hard_negatives

from torch.utils.data import DataLoader, Dataset
from torch.optim import Adam
from torch.optim.lr_scheduler import CosineAnnealingLR


class ContrastiveDataset(Dataset):
    """Dataset of (query_embedding, positive_doc_embedding) pairs."""

    def __init__(self, query_embeds, doc_embeds):
        self.query_embeds = torch.tensor(query_embeds, dtype=torch.float32)
        self.doc_embeds = torch.tensor(doc_embeds, dtype=torch.float32)

    def __len__(self):
        return len(self.query_embeds)

    def __getitem__(self, idx):
        return self.query_embeds[idx], self.doc_embeds[idx]


def train_contrastive_adapter(
    query_embeddings: np.ndarray,     # (2000, 1024)
    positive_embeddings: np.ndarray,  # (2000, 1024)
    catalog_embeddings: np.ndarray,   # (50000, 1024) — full catalog
    num_epochs: int = 20,
    batch_size: int = 128,
    learning_rate: float = 1e-4,
    temperature: float = 0.07,
):
    """
    Train contrastive adapter with hard negative mining.

    20 epochs is appropriate for our small (2K pairs) dataset.
    More epochs risk overfitting; fewer under-train the adapter.
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    adapter = ContrastiveAdapter(embed_dim=1024, bottleneck_dim=512).to(device)
    loss_fn = InfoNCELoss(temperature=temperature)
    optimizer = Adam(adapter.parameters(), lr=learning_rate, weight_decay=1e-5)
    scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs)

    miner = HardNegativeMiner(embed_dim=1024)

    best_recall = 0.0

    for epoch in range(num_epochs):
        # ── Hard Negative Refresh ──
        adapter.eval()
        with torch.no_grad():
            # Re-embed catalog through current adapter
            catalog_tensor = torch.tensor(catalog_embeddings, dtype=torch.float32).to(device)
            adapted_catalog = adapter(catalog_tensor).cpu().numpy()

        miner.build_index(adapted_catalog)
        # For now we use in-batch negatives; hard negative mining
        # enriches the batch composition by pre-selecting hard examples

        # ── Training ──
        dataset = ContrastiveDataset(query_embeddings, positive_embeddings)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

        adapter.train()
        epoch_loss = 0.0

        for query_batch, doc_batch in dataloader:
            query_batch = query_batch.to(device)
            doc_batch = doc_batch.to(device)

            # Forward: adapt both queries and documents
            adapted_queries = adapter(query_batch)
            adapted_docs = adapter(doc_batch)

            # Compute InfoNCE loss (in-batch negatives)
            loss = loss_fn(adapted_queries, adapted_docs)

            # Backward
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(adapter.parameters(), max_norm=1.0)
            optimizer.step()

            epoch_loss += loss.item()

        scheduler.step()

        # ── Validation ──
        recall_at_3 = evaluate_recall(adapter, device, query_embeddings, catalog_embeddings)
        avg_loss = epoch_loss / len(dataloader)
        print(f"Epoch {epoch+1}: Loss={avg_loss:.4f}, Recall@3={recall_at_3:.4f}")

        if recall_at_3 > best_recall:
            best_recall = recall_at_3
            torch.save(adapter.state_dict(), "best_adapter.pt")

    return adapter


def evaluate_recall(adapter, device, query_embeds, catalog_embeds, k=3):
    """Compute Recall@K on validation queries against full catalog."""
    adapter.eval()
    with torch.no_grad():
        q = torch.tensor(query_embeds, dtype=torch.float32).to(device)
        c = torch.tensor(catalog_embeds, dtype=torch.float32).to(device)
        q_adapted = adapter(q).cpu().numpy()
        c_adapted = adapter(c).cpu().numpy()

    # Build FAISS index
    index = faiss.IndexFlatIP(q_adapted.shape[1])
    index.add(c_adapted)

    # Search
    _, indices = index.search(q_adapted, k)

    # Compute recall (how often the true positive is in top-K)
    hits = 0
    for i, retrieved in enumerate(indices):
        if i in retrieved:  # Simplified: positive_id == query_id mapping
            hits += 1
    return hits / len(query_embeds)

import json
import torch
import numpy as np
import boto3

# ─── Cold-start optimization: load model at import time ───
adapter = ContrastiveAdapter(embed_dim=1024, bottleneck_dim=512)
adapter.load_state_dict(torch.load("/opt/ml/model/adapter.pt", map_location="cpu"))
adapter.eval()

bedrock = boto3.client("bedrock-runtime", region_name="us-east-1")


def lambda_handler(event, context):
    """
    Lambda function for real-time embedding adaptation.
    Called by the orchestrator between Titan V2 embedding and OpenSearch retrieval.

    P95 latency budget: 30ms total (Titan ~20ms + adapter ~2ms + overhead ~8ms)
    """
    query_text = event["query"]

    # Step 1: Get Titan Embedding (Bedrock API call)
    titan_response = bedrock.invoke_model(
        modelId="amazon.titan-embed-text-v2:0",
        body=json.dumps({
            "inputText": query_text,
            "dimensions": 1024,
            "normalize": True,
        }),
    )
    raw_embedding = json.loads(titan_response["body"].read())["embedding"]

    # Step 2: Adapt embedding through our trained adapter
    with torch.no_grad():
        input_tensor = torch.tensor([raw_embedding], dtype=torch.float32)
        adapted = adapter(input_tensor).squeeze(0).numpy()

    return {
        "statusCode": 200,
        "body": json.dumps({
            "embedding": adapted.tolist(),
            "dimension": 1024,
        }),
    }

graph LR
    subgraph "Training Pipeline (Bi-weekly)"
        A[Click-Through<br>Logs] --> B[Pair Extraction<br>+ Hard Neg Mining]
        B --> C[Adapter Training<br>Lambda GPU<br>(20 epochs, 18 min)]
        C --> D[Model Artifact<br>S3 (4MB)]
    end

    subgraph "Validation Gate"
        D --> E[Recall@3 > 0.80<br>on golden set]
        E -->|pass| F[Lambda Layer<br>Deployment]
        E -->|fail| G[Reject + Alert]
    end

    subgraph "Inference Pipeline"
        H[User Query] --> I[Titan V2<br>Bedrock API<br>~20ms]
        I --> J[Adapter Lambda<br>~2ms]
        J --> K[OpenSearch<br>ANN Search]
    end

    subgraph "Monitoring"
        K --> L[Recall@3 Tracking<br>(sampled)]
        L -->|drift detected| A
    end