MMR-Bench: Multimodal LLM Routing을 위한 종합 벤치마크

import numpy as np
from sklearn.cluster import KMeans
import clip
import torch

# CLIP으로 텍스트/이미지 임베딩 추출
model, preprocess = clip.load("ViT-B/32", device="cpu")

def get_embeddings(image, text):
    img_input = preprocess(image).unsqueeze(0)
    txt_input = clip.tokenize([text], truncate=True)
    with torch.no_grad():
        img_emb = model.encode_image(img_input).float()
        txt_emb = model.encode_text(txt_input).float()
    img_emb = img_emb / img_emb.norm(dim=-1, keepdim=True)
    txt_emb = txt_emb / txt_emb.norm(dim=-1, keepdim=True)
    return txt_emb.numpy()[0], img_emb.numpy()[0]

def adaptive_fuse(txt_emb, img_emb, eta=5.0, alpha=0.5, beta=0.5):
    """MMR-Bench 논문의 adaptive fusion (Eq. S6)"""
    # norm 기반 신뢰도로 softmax 가중치 계산
    c = np.array([np.linalg.norm(txt_emb), np.linalg.norm(img_emb)])
    weights = np.exp(eta * c) / np.exp(eta * c).sum()
    w_txt, w_img = weights

    # 가중합 + 곱(cross-modal agreement) + 차이(mismatch)
    fused = (w_txt * txt_emb + w_img * img_emb) \
          + alpha * (txt_emb * img_emb) \
          + beta  * (txt_emb - img_emb)
    return fused / np.linalg.norm(fused)

# 모델 zoo (비용: $/1M output tokens)
MODEL_ZOO = {
    "qwen-3b":      {"cost": 0.03},
    "qwen-7b":      {"cost": 0.07},
    "gemini-flash": {"cost": 2.00},
    "gpt-5":        {"cost": 10.0},
}

class KMeansRouter:
    def __init__(self, n_clusters=8):
        self.kmeans = KMeans(n_clusters=n_clusters, init='k-means++')
        self.cluster_model_map = {}

    def fit(self, embeddings, utilities, cost_weight=0.5):
        """utilities: (n_samples, n_models) 각 모델의 정답률"""
        self.kmeans.fit(embeddings)
        costs = np.array([v["cost"] for v in MODEL_ZOO.values()])
        for h in range(self.kmeans.n_clusters):
            mask = self.kmeans.labels_ == h
            scores = utilities[mask].mean(0) - cost_weight * (costs / costs.max())
            self.cluster_model_map[h] = list(MODEL_ZOO.keys())[np.argmax(scores)]

    def route(self, embedding):
        cluster = self.kmeans.predict(embedding.reshape(1, -1))[0]
        return self.cluster_model_map[cluster]

# --- 추론 시 ---
# txt_emb, img_emb = get_embeddings(image, query)
# fused = adaptive_fuse(txt_emb, img_emb)
# selected = router.route(fused)
# print(f"→ {selected} (${MODEL_ZOO[selected]['cost']}/1M tokens)")

# 텍스트 전용 쿼리라면 이미지 임베딩을 0으로 마스킹 (재훈련 불필요)
# fused = adaptive_fuse(txt_emb, np.zeros_like(txt_emb))