Don't Start Over: A Cost-Effective Framework for Migrating Personalized Prompts Between LLMs

# PUMA adapter structure (PyTorch-like code)
import torch
import torch.nn as nn

class PUMAAdapter(nn.Module):
    """
    source_dim: old LLM embedding dimension (e.g., 1B model)
    target_dim: new LLM embedding dimension (e.g., 3B model)
    prompt_len: soft prompt length (l=1 in the paper)
    """
    def __init__(self, source_dim: int, target_dim: int):
        super().__init__()
        self.adapter = nn.Sequential(
            nn.Linear(source_dim, target_dim * 2),
            nn.LayerNorm(target_dim * 2),
            nn.GELU(),
            nn.Linear(target_dim * 2, target_dim),
        )
        # Projection for residual connection
        self.residual_proj = nn.Linear(source_dim, target_dim)

    def forward(self, source_prompt: torch.Tensor) -> torch.Tensor:
        # source_prompt: (batch, prompt_len, source_dim)
        return self.adapter(source_prompt) + self.residual_proj(source_prompt)

# User selection strategy (group-based)
from sklearn.cluster import KMeans
import numpy as np

def select_representative_users(
    prompt_embeddings: np.ndarray,  # (num_users, emb_dim)
    output_variance: np.ndarray,    # (num_users,)
    n_clusters: int = 50,
    budget: int = 2000,
) -> list[int]:
    # Stage 1: Preference diversity clustering with K-means
    kmeans = KMeans(n_clusters=n_clusters, random_state=42)
    cluster_labels = kmeans.fit_predict(prompt_embeddings)

    selected_indices = []
    per_cluster_budget = budget // n_clusters

    for c in range(n_clusters):
        cluster_mask = cluster_labels == c
        cluster_idx = np.where(cluster_mask)[0]
        cluster_var = output_variance[cluster_idx]

        # Stage 2: Variance-based stratified sampling (weight toward mid-variance users)
        bins = np.percentile(cluster_var, [33, 66])
        low = cluster_idx[cluster_var <= bins[0]]
        mid = cluster_idx[(cluster_var > bins[0]) & (cluster_var <= bins[1])]
        high = cluster_idx[cluster_var > bins[1]]

        # Normal distribution weights: allocate more to the middle group
        weights = [1, 2, 1]  # low:mid:high
        total_w = sum(weights)
        for group, w in zip([low, mid, high], weights):
            n = max(1, int(per_cluster_budget * w / total_w))
            if len(group) > 0:
                chosen = np.random.choice(group, min(n, len(group)), replace=False)
                selected_indices.extend(chosen.tolist())

    return selected_indices[:budget]