임베딩 모델의 테스트 시 컴퓨팅 스케일링

OpenAI가 O1 모델을 공개한 이후，AI 커뮤니티에서 가장 많이 논의된 주제 중 하나는 test-time compute 확장입니다. 이는 사전 학습 단계가 아닌 추론—AI 모델이 입력에 대한 출력을 생성하는 단계—중에 추가적인 컴퓨팅 리소스를 할당하는 것을 의미합니다. 잘 알려진 예시로는 "사고의 연쇄" 다단계 추론이 있는데，이를 통해 모델은 여러 가능한 답변을 평가하고，더 깊이 있는 계획을 세우며，최종 응답을 하기 전에 자기 성찰을 하는 등 더 광범위한 내부 심의를 수행할 수 있습니다. 이 전략은 특히 복잡한 추론 작업에서 답변의 질을 향상시킵니다. 최근 Alibaba가 공개한 QwQ-32B-Preview 모델도 test-time compute 증가를 통한 AI 추론 개선이라는 이러한 트렌드를 따르고 있습니다.

💡

이 맥락에서 "확장"은 주로 추론 중에 사용 가능한 컴퓨팅 용량(처리 능력이나 시간 등)을 늘리는 것을 의미합니다. 이는 scaling out(여러 시스템에 걸쳐 작업 분산) 또는 speedup(처리 시간 단축)을 의미하지 않습니다.

0:00

/0:10

OpenAI의 O1 모델을 사용할 때，모델이 문제 해결을 위한 추론 체인을 구성하면서 다단계 추론에 추가 시간이 필요하다는 것을 사용자들이 명확히 인식할 수 있습니다.

Jina AI에서는 LLM보다는 임베딩과 리랭커에 더 초점을 맞추고 있기 때문에，이러한 맥락에서 test-time compute 확장을 고려하는 것이 자연스럽습니다: "사고의 연쇄"를 임베딩 모델에 어떻게 적용할 수 있을까요? 처음에는 직관적으로 보이지 않을 수 있지만，이 글에서는 새로운 관점을 탐구하고 test-time compute 확장을 jina-clip에 적용하여 분포 외(OOD) 이미지를 분류하는—그렇지 않으면 불가능할 작업을 해결하는 방법을 보여줍니다.

우리의 실험은 임베딩 모델에게 흥미로운 도전이 되는 포켓몬 인식에 초점을 맞췄습니다. CLIP과 같은 모델들이 일반적인 이미지-텍스트 매칭에서 뛰어난 성능을 보이지만，미세 조정 없이는 틈새 도메인이나 OOD 이미지에서 어려움을 겪을 수 있습니다. 모델에게 더 많은 "생각할 시간"을 주면，"사고의 연쇄"와 유사한 다중 대상 분류가 임베딩 모델 자체의 조정 없이도 정확도를 향상시킬 수 있다는 것을 발견했습니다.

tag사례 연구

TheFusion21/PokemonCards 데이터셋을 사용한 포켓몬 분류에 초점을 맞춘 우리의 실험은 수천 개의 포켓몬 트레이딩 카드 이미지를 포함하고 있습니다. 이미지 분류 작업에서 입력은 잘라낸 포켓몬 카드 아트워크(모든 텍스트/설명 제거)이며 출력은 미리 정의된 이름 목록에서 정확한 포켓몬 이름입니다. 이 작업은 CLIP 임베딩 모델에게 특히 흥미로운 도전이 되는데，그 이유는:

포켓몬 이름과 시각적 특징이 모델에게는 틈새의，분포 외 개념을 나타내므로 직접적인 분류가 어려움
각 포켓몬은 CLIP이 더 잘 이해할 수 있는 기본 요소(모양，색상，자세)로 분해할 수 있는 명확한 시각적 특징을 가짐
카드 아트워크는 일관된 시각적 형식을 제공하면서도 다양한 배경，자세，예술적 스타일을 통해 복잡성을 도입
이 작업은 언어 모델의 복잡한 추론 체인과 유사하게 여러 시각적 특징을 동시에 통합해야 함

해당 텍스트에 포켓몬 이름이 나타나는 것으로 인한 단순 추측을 방지하기 위해 모든 텍스트 정보(헤더，푸터，설명)를 제거하여 포켓몬 카드 이미지를 잘라냈습니다. 이 포켓몬들의 클래스 레이블은 [`Absol G`，`Aerodactyl`，`Weedle`，`Caterpie`, `Azumarill`, `Bulbasaur`, `Venusaur`, `Absol`, `Aggron`, `Beedrill δ`, `Alakazam`, `Ampharos`, `Dratini`, `Ampharos`, `Ampharos`, `Arcanine`, `Blaine's Moltres`, `Aerodactyl`, `Celebi & Venusaur-GX`, `Caterpie`]

tag기준선

기준선 접근 방식은 포켓몬 카드 아트워크와 이름 간의 단순 직접 비교를 사용합니다. 먼저 각 포켓몬 카드 이미지에서 모든 텍스트 정보(헤더, 푸터, 설명)를 제거하여 CLIP 모델이 해당 텍스트에 나타나는 포켓몬 이름으로 인한 간단한 추측을 하지 못하도록 합니다. 그런 다음 jina-clip-v1과 jina-clip-v2 모델을 사용하여 잘라낸 이미지와 포켓몬 이름을 인코딩하여 각각의 임베딩을 얻습니다. 분류는 이러한 이미지와 텍스트 임베딩 간의 코사인 유사도를 계산하여 수행됩니다 - 각 이미지는 가장 높은 유사도 점수를 가진 이름과 매칭됩니다. 이는 추가적인 맥락이나 속성 정보 없이 카드 아트워크와 포켓몬 이름 간의 일대일 매칭을 생성합니다. 아래의 의사 코드는 기준선 방법을 요약합니다.

# Preprocessing
cropped_images = [crop_artwork(img) for img in pokemon_cards]  # Remove text, keep only art
pokemon_names = ["Absol", "Aerodactyl", ...]  # Raw Pokemon names

# Get embeddings using jina-clip-v1
image_embeddings = model.encode_image(cropped_images)
text_embeddings = model.encode_text(pokemon_names)

# Classification by cosine similarity
similarities = cosine_similarity(image_embeddings, text_embeddings)
predicted_names = [pokemon_names[argmax(sim)] for sim in similarities]

# Evaluate
accuracy = mean(predicted_names == ground_truth_names)

tag분류를 위한 "사고의 연쇄"

이미지를 이름과 직접 매칭하는 대신, 우리는 포켓몬 인식을 시각적 속성의 구조화된 시스템으로 분해합니다. 우리는 다섯 가지 주요 속성 그룹을 정의합니다: 지배적인 색상(예: "흰색", "파란색"), 기본 형태(예: "늑대", "날개 달린 파충류"), 주요 특징(예: "하나의 흰색 뿔", "큰 날개"), 몸체 형태(예: "네 발로 선 늑대 같은", "날개 달린 날씬한"), 배경 장면(예: "우주", "녹색 숲").

각 속성 그룹에 대해, 우리는 특정 텍스트 프롬프트(예: "이 포켓몬의 몸은 주로 {} 색상입니다")와 관련 옵션들을 생성합니다. 그런 다음 모델을 사용하여 이미지와 각 속성 옵션 간의 유사도 점수를 계산합니다. 이러한 점수들은 소프트맥스를 사용하여 확률로 변환되어 더 보정된 신뢰도 측정값을 얻습니다.

전체 사고의 연쇄(CoT) 구조는 두 부분으로 구성됩니다: 프롬프트 그룹을 설명하는 classification_groups와 각 포켓몬이 매칭해야 할 속성 옵션을 정의하는 pokemon_rules입니다. 예를 들어, Absol은 색상에서 "흰색"과 형태에서 "늑대 같은"과 매칭되어야 합니다. 전체 CoT는 아래와 같습니다(이것이 어떻게 구성되는지는 나중에 설명하겠습니다):

pokemon_system = {
    "classification_cot": {
        "dominant_color": {
            "prompt": "이 포켓몬의 몸은 주로 {} 색상입니다.",
            "options": [
                "white",    # Absol, Absol G
                "gray",     # Aggron
                "brown",    # Aerodactyl, Weedle, Beedrill δ
                "blue",     # Azumarill
                "green",    # Bulbasaur, Venusaur, Celebi&Venu, Caterpie
                "yellow",   # Alakazam, Ampharos
                "red",      # Blaine's Moltres
                "orange",   # Arcanine
                "light blue"# Dratini
            ]
        },
        "primary_form": {
            "prompt": "이것은 {} 처럼 보입니다.",
            "options": [
                "a wolf",         # Absol, Absol G
                "an armored dinosaur",  # Aggron
                "a winged reptile",     # Aerodactyl
                "a rabbit-like creature", # Azumarill
                "a toad-like creature",   # Bulbasaur, Venusaur, Celebi&Venu
                "a caterpillar larva",    # Weedle, Caterpie
                "a wasp-like insect",     # Beedrill δ
                "a fox-like humanoid",     # Alakazam
                "a sheep-like biped",      # Ampharos
                "a dog-like beast",        # Arcanine
                "a flaming bird",          # Blaine's Moltres
                "a serpentine dragon"      # Dratini
            ]
        },
        "key_trait": {
            "prompt": "가장 눈에 띄는 특징은 {} 입니다.",
            "options": [
                "a single white horn", # Absol, Absol G
                "metal armor plates",  # Aggron
                "large wings",         # Aerodactyl, Beedrill δ
                "rabbit ears",         # Azumarill
                "a green plant bulb",  # Bulbasaur, Venusaur, Celebi&Venu
                "a small red spike",   # Weedle
                "big green eyes",      # Caterpie
                "a mustache and spoons", # Alakazam
                "a glowing tail orb",  # Ampharos
                "a fiery mane",        # Arcanine
                "flaming wings",       # Blaine's Moltres
                "a tiny white horn on head" # Dratini
            ]
        },
        "body_shape": {
            "prompt": "몸체 형태는 {} 로 묘사될 수 있습니다.",
            "options": [
                "wolf-like on four legs",   # Absol, Absol G
                "bulky and armored",        # Aggron
                "winged and slender",       # Aerodactyl, Beedrill δ
                "round and plump",          # Azumarill
                "sturdy and four-legged",   # Bulbasaur, Venusaur, Celebi&Venu
                "long and worm-like",       # Weedle, Caterpie
                "upright and humanoid",     # Alakazam, Ampharos
                "furry and canine",         # Arcanine
                "bird-like with flames",    # Blaine's Moltres
                "serpentine"                # Dratini
            ]
        },
        "background_scene": {
            "prompt": "배경은 {} 처럼 보입니다.",
            "options": [
                "outer space",      # Absol G, Beedrill δ
                "green forest",     # Azumarill, Bulbasaur, Venusaur, Weedle, Caterpie, Celebi&Venu
                "a rocky battlefield", # Absol, Aggron, Aerodactyl
                "a purple psychic room", # Alakazam
                "a sunny field",     # Ampharos
                "volcanic ground",   # Arcanine
                "a red sky with embers", # Blaine's Moltres
                "a calm blue lake"   # Dratini
            ]
        }
    },
    
    "pokemon_rules": {
        "Absol": {
            "dominant_color": 0,      
            "primary_form": 0,   
            "key_trait": 0,      
            "body_shape": 0,    
            "background_scene": 2   
        },
        "Absol G": {
            "dominant_color": 0,      
            "primary_form": 0,   
            "key_trait": 0,       
            "body_shape": 0,     
            "background_scene": 0    
        },
        // ...
    }
}

최종 분류는 이러한 속성 확률들을 결합합니다 - 단일 유사도 비교 대신, 이제 우리는 여러 구조화된 비교를 수행하고 그들의 확률을 집계하여 더 정보에 기반한 결정을 내립니다.

# Classification process
def classify_pokemon(image):
   # Generate all text prompts
   all_prompts = []
   for group in classification_cot:
       for option in group["options"]:
           prompt = group["prompt"].format(option)
           all_prompts.append(prompt)

   # Get embeddings and similarities
   image_embedding = model.encode_image(image)
   text_embeddings = model.encode_text(all_prompts)
   similarities = cosine_similarity(image_embedding, text_embeddings)

   # Convert to probabilities per attribute group
   probabilities = {}
   for group_name, group_sims in group_similarities:
       probabilities[group_name] = softmax(group_sims)

   # Score each Pokemon based on matching attributes
   scores = {}
   for pokemon, rules in pokemon_rules.items():
       score = 0
       for group, target_idx in rules.items():
           score += probabilities[group][target_idx]
       scores[pokemon] = score

   return max(scores, key=scores.get)

tag복잡도 분석

이미지를 N개의 포켓몬 이름 중 하나로 분류하고자 한다고 가정해봅시다. 기준선 접근 방식은 N개의 텍스트 임베딩(각 포켓몬 이름당 하나)을 계산해야 합니다. 반면에 우리의 스케일된 테스트-타임 계산 접근 방식은 Q개의 텍스트 임베딩을 계산해야 하며, 여기서Q는 모든 질문에 걸친 질문-옵션 조합의 총 개수입니다. 두 방법 모두 하나의 이미지 임베딩을 계산하고 최종 분류 단계를 수행해야 하므로, 이러한 공통 연산은 비교에서 제외합니다. 이 사례 연구에서는 N=13이고 Q=52입니다.

Q = N인 극단적인 경우, 우리의 접근 방식은 기본적으로 기준선과 동일해질 것입니다. 하지만 테스트 시간 계산을 효과적으로 확장하는 핵심은 다음과 같습니다:

Q를 증가시키는 신중하게 선택된 질문들을 구성
각 질문이 최종 답변에 대한 명확하고 유익한 단서를 제공하도록 보장
질문들을 최대한 직교하도록 설계하여 공동 정보 이득을 최대화

이 접근 방식은 각 질문이 가능한 답변을 효과적으로 좁히기 위해 전략적으로 선택되는 "스무고개" 게임과 유사합니다.

tag평가

평가는 13개의 서로 다른 포켓몬 클래스에 걸쳐 117개의 테스트 이미지로 수행되었습니다. 결과는 다음과 같습니다:

Approach	jina-clip-v1	jina-clip-v2
Baseline	31.36%	16.10%
CoT	46.61%	38.14%
Improvement	+15.25%	+22.04%

동일한 CoT 분류가 이러한 일반적이지 않은 또는 OOD 작업에서 두 모델 모두에 대해 상당한 개선(각각 +15.25%와 +22.04%)을 제공하는 것을 볼 수 있습니다. 이는 또한 pokemon_system이 구축되면, 동일한 CoT 시스템을 다른 모델들에 효과적으로 전이할 수 있으며, 미세 조정이나 사후 학습이 필요하지 않다는 것을 시사합니다.

v1의 상대적으로 강력한 기준선 성능(31.36%)은 주목할 만합니다. 이 모델은 포켓몬 관련 콘텐츠를 포함한 LAION-400M으로 학습되었습니다. 반면에 v2는 더 높은 품질이지만 더 필터링된 데이터셋인 DFN-2B(400M 인스턴스 서브샘플링)에서 학습되었는데, 이는 포켓몬 관련 콘텐츠를 제외했을 수 있어 이 특정 작업에서 v2의 더 낮은 기준선 성능(16.10%)을 설명합니다.

tag`pokemon_system` 효과적으로 구축하기

우리의 확장된 테스트 시간 계산 접근 방식의 효과는 pokemon_system을 얼마나 잘 구축하는지에 크게 달려있습니다. 이 시스템을 구축하는 데는 수동부터 완전 자동화까지 다양한 접근 방식이 있습니다.

수동 구축

가장 직접적인 접근 방식은 포켓몬 데이터셋을 수동으로 분석하고 속성 그룹, 프롬프트, 규칙을 만드는 것입니다. 도메인 전문가는 색상, 형태, 특징적인 특성과 같은 주요 시각적 속성을 식별해야 합니다. 그런 다음 각 속성에 대한 자연어 프롬프트를 작성하고, 각 속성 그룹에 대한 가능한 옵션을 열거하며, 각 포켓몬을 올바른 속성 옵션에 매핑해야 합니다. 이는 고품질 규칙을 제공하지만, 시간이 많이 소요되며 더 큰 N에 대해 잘 확장되지 않습니다.

LLM 지원 구축

우리는 LLM에 프롬프트를 제공하여 분류 시스템을 생성하도록 함으로써 이 프로세스를 가속화할 수 있습니다. 잘 구조화된 프롬프트는 시각적 특성을 기반으로 한 속성 그룹, 자연어 프롬프트 템플릿, 포괄적이고 상호 배타적인 옵션, 각 포켓몬에 대한 매핑 규칙을 요청할 것입니다. LLM은 초안을 빠르게 생성할 수 있지만, 그 출력은 검증이 필요할 수 있습니다.

I need help creating a structured system for Pokemon classification. For each Pokemon in this list: [Absol, Aerodactyl, Weedle, Caterpie, Azumarill, ...], create a classification system with:

1. Classification groups that cover these visual attributes:
   - Dominant color of the Pokemon
   - What type of creature it appears to be (primary form)
   - Its most distinctive visual feature
   - Overall body shape
   - What kind of background/environment it's typically shown in

2. For each group:
   - Create a natural language prompt template using "{}" for the option
   - List all possible options that could apply to these Pokemon
   - Make sure options are mutually exclusive and comprehensive

3. Create rules that map each Pokemon to exactly one option per attribute group, using indices to reference the options

Please output this as a Python dictionary with two main components:
- "classification_groups": containing prompts and options for each attribute
- "pokemon_rules": mapping each Pokemon to its correct attribute indices

Example format:
{
    "classification_groups": {
        "dominant_color": {
            "prompt": "This Pokemon's body is mainly {} in color",
            "options": ["white", "gray", ...]
        },
        ...
    },
    "pokemon_rules": {
        "Absol": {
            "dominant_color": 0,  # index for "white"
            ...
        },
        ...
    }
}

더 강력한 접근 방식은 LLM 생성과 인간 검증을 결합하는 것입니다. 먼저 LLM이 초기 시스템을 생성합니다. 그런 다음 전문가들이 속성 그룹화, 옵션의 완전성, 규칙의 정확성을 검토하고 수정합니다. LLM은 이 피드백을 바탕으로 시스템을 개선하며, 만족스러운 품질이 달성될 때까지 이 과정이 반복됩니다. 이 접근 방식은 효율성과 정확성의 균형을 맞춥니다.

DSPy를 통한 자동화된 구축

완전 자동화된 접근 방식의 경우, DSPy를 사용하여 pokemon_system을 반복적으로 최적화할 수 있습니다. 이 과정은 수동으로 또는 LLM에 의해 작성된 간단한 pokemon_system을 초기 프롬프트로 시작합니다. 각 버전은 홀드아웃 세트에서 평가되며, 정확도를 DSPy의 피드백 신호로 사용합니다. 이 성능을 기반으로 최적화된 프롬프트(즉, pokemon_system의 새 버전)가 생성됩니다. 이 사이클은 수렴할 때까지 반복되며, 전체 과정 동안 임베딩 모델은 완전히 고정된 상태로 유지됩니다.

DSPy를 사용하여 최적의 `pokemon_system` CoT 설계를 찾기; 튜닝 프로세스는 각 작업에 대해 한 번만 수행하면 됨.

tag임베딩 모델에 대한 테스트 시간 계산을 확장하는 이유는?

사전 학습의 확장이 결국 경제적으로 감당할 수 없게 되기 때문입니다.

Jina 임베딩 제품군(jina-embeddings-v1, v2, v3, jina-clip-v1, v2, jina-ColBERT-v1, v2 포함) 출시 이후, 확장된 사전 학습을 통한 각 모델 업그레이드에는 더 많은 비용이 수반되었습니다. 예를 들어, 2023년 6월에 출시된 첫 번째 모델인 jina-embeddings-v1은 110M 매개변수를 가졌습니다. 당시 학습 비용은 측정 방식에 따라 $5,000에서$ 10,000 사이였습니다. jina-embeddings-v3의 개선은 상당하지만, 주로 투자된 리소스 증가에서 비롯됩니다. 최신 모델의 비용 궤적은 수천 달러에서 수만 달러로, 더 큰 AI 기업의 경우 현재 수억 달러까지 증가했습니다. 사전 학습에 더 많은 돈, 리소스, 데이터를 투입하면 더 나은 모델이 생성되지만, 한계 수익은 결국 추가 확장을 경제적으로 지속 불가능하게 만듭니다.

임베딩 모델의 확장 법칙. 영어 작업에 대한 평균 MTEB 성능이 모델 매개변수 수에 대해 플롯되었습니다. 각 점은 임베딩 모델을 나타냅니다. 모든 모델을 대표하는 추세선이 강조되었고, 다국어 모델은 청록색 점으로 표시됩니다. 이 그래프는 MTEB 리더보드에서 상위 100개 임베딩 모델을 선택하여 작성되었으며, 크기 정보가 없는 모델(일반적으로 비공개 또는 독점 모델)은 제외되었습니다. 명백한 트롤링으로 확인된 제출물도 필터링되었습니다.

반면에 현대의 임베딩 모델들은 점점 더 강력해지고 있습니다: 다국어, 멀티태스크, 멀티모달, 그리고 강력한 제로샷 및 지시 따르기 성능을 갖추고 있습니다. 이러한 다재다능함은 알고리즘 개선과 테스트 시간 계산 확장을 위한 큰 여지를 남겨둡니다.

그렇다면 질문은 이렇게 됩니다: 사용자들이 깊이 신경 쓰는 쿼리에 대해 기꺼이 지불할 비용은 얼마일까요? 고정된 사전 학습 모델에 대해 더 긴 추론 시간을 감수하면서 결과의 품질을 크게 개선할 수 있다면, 많은 사람들이 그것을 가치 있게 여길 것입니다. 우리의 관점에서는 임베딩 모델에 대한 테스트 시간 계산 확장에 상당한 미개발 잠재력이 있습니다. 이는 단순히 학습 중 모델 크기를 증가시키는 것에서 더 나은 성능을 달성하기 위해 추론 단계에서 계산 노력을 강화하는 방향으로의 전환을 나타냅니다.

tag결론

jina-clip-v1/v2의 테스트 시간 계산에 대한 우리의 사례 연구는 몇 가지 주요 발견을 보여줍니다:

임베딩에 대한 미세 조정이나 사후 학습 없이도 일반적이지 않은 또는 분포 외(OOD) 데이터에서 더 나은 성능을 달성했습니다.
시스템은 유사도 검색과 분류 기준을 반복적으로 개선함으로써 더 미묘한 구분을 만들었습니다.
동적 프롬프트 조정과 반복적 추론을 통합함으로써, 임베딩 모델의 추론 과정을 단일 쿼리에서 더 정교한 사고 체인으로 변환했습니다.

이 사례 연구는 테스트 시간 계산으로 가능한 것의 표면만을 긁었을 뿐입니다. 알고리즘적으로 확장할 수 있는 상당한 여지가 남아있습니다. 예를 들어, "스무고개" 게임의 최적 전략과 유사하게 답변 공간을 가장 효율적으로 좁힐 수 있는 질문을 반복적으로 선택하는 방법을 개발할 수 있습니다. 테스트 시간 계산을 확장함으로써, 우리는 임베딩 모델을 현재의 한계를 넘어 밀어붙여 한때 도달할 수 없어 보였던 더 복잡하고 미묘한 작업을 해결할 수 있게 할 수 있습니다.