Jina 코드 임베딩: 0.5B 및 1.5B에서 SOTA 코드 검색

Efficient Code Embeddings from Code Generation Models

jina-code-embeddings는 자연어 쿼리에서 코드를 검색하고, 기술적인 질의응답을 수행하며, 프로그래밍 언어 전반에서 의미상 유사한 코드 스니펫을 식별하도록 설계된 새로운 코드 向量模型 제품군입니다. 텍스트와 코드 모두에서 사전 훈련된 자기 회귀 백본을 혁신적으로 사용하여 마지막 词元 풀링을 통해 向量模型을 생성합니다. 모델의 상대적으로 작은 크기에도 불구하고 최첨단 성능을 입증하는 훈련 레시피를 간략하게 설명하고 코드 向量模型 구성에 대한 이 접근 방식을 검증합니다.

arXiv.orgDaria Kryvosheieva

오늘 저희는 두 가지 크기(0.5B 및 1.5B 파라미터)의 새로운 코드 向量模型 제품군인 jina-code-embeddings와 두 모델 모두에 대한 GGUF 양자화를 출시합니다. 자동 회귀 코드 생성 LLM을 기반으로 구축된 이 모델은 컴팩트한 크기에도 불구하고 최첨단 검색 성능을 달성합니다. 이 모델은 Python, JavaScript, Java, C++, C#, Go, Rust, TypeScript, SQL, MATLAB, R, Swift, Kotlin, HTML/CSS, PHP, Ruby, Scala, Perl 및 Shell을 포함한 15개 이상의 프로그래밍 언어를 지원합니다.

jina-code-embeddings는 25개의 코드 검색 벤치마크에서 78.41%(0.5B) 및 79.04%(1.5B)의 평균 성능을 달성합니다. 0.5B 모델은 Qwen3-Embedding-0.6B보다 20% 더 작음에도 불구하고 5%p 더 높은 성능을 보이며, 1.5B 변형은 voyage-code-3(79.23%)와 유사하고 아키텍처가 공개되지 않은 독점 모델인 gemini-embedding-001(77.38%)을 능가합니다.

모델	파라미터	전체 평균	MTEB 코드 평균
<strong>jina-code-embeddings-1.5b</strong>	1.54B	79.04%	78.94%
<strong>jina-code-embeddings-0.5b</strong>	494M	78.41%	78.72%
voyage-code-3	알 수 없음*	79.23%	79.84%
gemini-embedding-001	알 수 없음*	77.38%	76.48%
jina-embeddings-v4	3.8B	74.11%	74.87%
Qwen3-Embedding-0.6B	600M	73.49%	74.69%

*아키텍처가 공개되지 않은 폐쇄 소스 모델

두 모델 모두 비대칭 검색을 위해 쿼리 및 문서 역할을 모두 지원하는 다양한 검색 시나리오에 대한 5개의 작업별 명령어 접두사로 훈련되었습니다. 예를 들어 nl2code_query를 사용하여 쿼리를 포함하고 nl2code_document를 사용하여 문서를 포함할 수 있습니다.

작업	사용 사례	명령어 접두사
`nl2code`	"CSV 읽는 방법" → `pandas.read_csv()`	"다음 쿼리가 주어졌을 때 가장 관련성이 높은 코드 스니펫을 찾으세요:\n"
`qa`	기술 Q&A 검색	"다음 질문이 주어졌을 때 가장 관련성이 높은 답변을 찾으세요:\n"
`code2code`	유사한 구현 찾기	"다음 코드 스니펫이 주어졌을 때 동등한 코드 스니펫을 찾으세요:\n"
`code2nl`	코드에서 문서로	"다음 코드 스니펫이 주어졌을 때 가장 관련성이 높은 주석을 찾으세요:\n"
`code2completion`	자동 완성 시나리오	"다음 코드 스니펫의 시작 부분이 주어졌을 때 가장 관련성이 높은 완성을 찾으세요:\n"

tag훈련 레시피

저희는 사전 훈련된 코드 생성 모델을 向量模型 백본으로 사용합니다. Qwen2.5-Coder-0.5B 및 1.5B를 기반으로 구축된 저희 모델은 다음과 같은 특징을 가지고 있습니다.

기능	jina-code-embeddings-0.5b	jina-code-embeddings-1.5b
기본 모델	Qwen2.5-Coder-0.5B	Qwen2.5-Coder-1.5B
向量模型 차원	896	1536
마트료시카 차원	64, 128, 256, 512, 896	128, 256, 512, 1024, 1536
최대 시퀀스 길이	32,768 词元	32,768 词元
풀링 전략	마지막 词元 풀링	마지막 词元 풀링
어텐션	FlashAttention2	FlashAttention2
데이터 유형	BFloat16	BFloat16

기존 코드 向量模型은 근본적인 병목 현상에 직면해 있습니다. 즉, 감독 훈련을 위한 고품질 주석-코드 쌍이 충분하지 않습니다. 92개 이상의 프로그래밍 언어에 걸쳐 5조 5천억 개의 词元으로 사전 훈련된 Qwen2.5-Coder로 시작함으로써 프로그래밍 구조, 교차 언어 패턴 인식 및 구문 및 관용구에 대한 기본 제공 지식에 대한 심층적인 의미론적 이해를 상속받습니다. 그런 다음 대조 학습 미세 조정은 최소한의 정렬된 데이터를 사용하여 검색 작업에 이 지식을 적용하여 인코더 전용 모델을 제약하는 데이터 부족을 회피합니다.

교차 프레임워크 코드 변환과 같이 과소 대표되는 작업의 경우 LLM을 사용하여 합성 데이터를 생성했으며 모든 합성 예제의 품질을 수동으로 검증했습니다. 저희의 훈련 데이터는 기존 MTEB 코드 작업 훈련 분할을 CommitPackFT, SWE-Bench, Spider, MBPP 및 CodeSearchNet을 포함한 조정된 공개 데이터 세트와 결합했습니다.

jina-embeddings-v3 및 v4와 달리 LoRA를 사용하지 않고 곧바로 전체 사후 훈련으로 진행했습니다. 저희와 같은 작은 모델(494M 및 1.54B 파라미터)의 경우 LoRA의 파라미터 효율성은 덜 매력적입니다. 어댑터 오버헤드는 용량이 제한적인 경우 실제로 성능을 저하시킬 수 있습니다. 저희는 모든 파라미터가 向量模型 작업에 사용되도록 해야 했습니다. 다중 작업 시나리오의 경우에도 작업별 명령어 접두사가 여러 LoRA 어댑터보다 더 깔끔하다는 것이 입증되었습니다. 가중치 구성을 전환하는 대신 단순히 다른 명령어를 추가합니다. 이는 LLM이 조건부 정보를 자연스럽게 처리하는 방식과 훨씬 더 간결하고 일치합니다.

훈련은 놀라울 정도로 효율적이었습니다. 두 모델 모두 4x A100 80GB GPU에서 InfoNCE 손실을 사용하여 대조 학습을 사용하여 훈련되었으며 0.5B 모델의 경우 8.3시간, 1.5B 변형의 경우 12시간 만에 완료되었습니다.

마지막으로 다양한 풀링 전략을 벤치마킹했습니다. 마지막 词元 풀링은 78.41%의 전체 평균을 달성하여 모든 벤치마크 범주에서 평균 풀링(77.20%) 및 잠재 어텐션 풀링(78.27%)보다 꾸준히 더 나은 성능을 보였습니다. 이 1.2%p의 이점은 jina-embeddings-v2, v3 및 v4에서 확립한 평균 풀링 전통에서 벗어나게 했습니다. 더 많은 검색 모델이 디코더 전용 LLM을 기반으로 구축됨에 따라 마지막 词元 풀링이 자연스러운 선택이 되었습니다. 평균 풀링은 단방향 어텐션 메커니즘과 잘 맞지 않습니다. 평균 풀링이 작동할 수 있고 초기 단계에서 더 쉽게 훈련될 수 있지만(볼록 최적화 환경 때문일 가능성이 높음) 저희의 실험에서는 마지막 词元 풀링이 달성하는 성능 한계 아래에서 평준화된다는 것을 일관되게 보여줍니다.

tag시작하기

두 모델 모두 Search Foundation API와 sentence-transformers, transformers 및 llama.cpp를 포함한 인기 있는 프레임워크를 통해 원활하게 작동합니다.

tagAPI를 통해

curl http://api.jina.ai/v1/embeddings \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $JINA_API_KEY" \
  -d @- <<EOFEOF
  {
    "model": "jina-code-embeddings-1.5b",
    "input": ["print hello world in python"],
    "task": "nl2code.passage"
  }
EOFEOF

tag`sentence-transformers`를 통해

from sentence_transformers import SentenceTransformer

# 모델 로드 (0.5b 또는 1.5b 선택)
model = SentenceTransformer(
    "jinaai/jina-code-embeddings-1.5b",
    model_kwargs={"torch_dtype": "bfloat16"},
    tokenizer_kwargs={"padding_side": "left"}
)

# 자연어에서 코드로
queries = ["print hello world in python", "initialize array of 5 zeros in c++"]
documents = ["print('Hello World!')", "int arr[5] = {0, 0, 0, 0, 0};"]

# 작업별 접두사로 向量模型 생성
query_embeddings = model.encode(queries, prompt_name="nl2code_query")
document_embeddings = model.encode(documents, prompt_name="nl2code_document")

# 유사성 계산
similarity = model.similarity(query_embeddings, document_embeddings)

tag`transformers`를 통해

from transformers import AutoModel, AutoTokenizer
import torch.nn.functional as F

def last_token_pool(last_hidden_states, attention_mask):
    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
    if left_padding:
        return last_hidden_states[:, -1]
    else:
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = last_hidden_states.shape[0]
        return last_hidden_states[torch.arange(batch_size), sequence_lengths]

tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-code-embeddings-1.5b')
model = AutoModel.from_pretrained('jinaai/jina-code-embeddings-1.5b')

# Apply task-specific prefix
query = "Find the most relevant code snippet given the following query:\nprint hello world"
code = "Candidate code snippet:\nprint('Hello World!')"

# Tokenize and embed
batch_dict = tokenizer([query, code], padding=True, truncation=True, return_tensors="pt")
outputs = model(**batch_dict)
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

tagMatryoshka 向量모델 컷오프

두 모델 모두 차원 [64, 128, 256, 512, 896]에 대해 Matryoshka 표현 학습으로 훈련되었으므로 다시 계산하지 않고도 向量모델을 잘라낼 수 있습니다.

# Full embeddings: 896d (0.5B) or 1536d (1.5B)
full_embedding = model.encode(text)

# Truncate to smaller dimensions for efficiency
small_embedding = full_embedding[:256]  # Works for both models
tiny_embedding = full_embedding[:128]   # 0.5B supports down to 64d

이러한 유연성을 통해 요구 사항에 따라 성능과 효율성 사이에서 균형을 맞출 수 있습니다.

tag결론

jina-code-embeddings는 효과적인 코드 向量모델이 대규모일 필요가 없다는 것을 보여줍니다. 코드 생성 모델을 기반으로 구축하고 대상에 맞는 미세 조정을 적용하여 15억 개 미만의 매개변수를 가진 모델로 최첨단 성능을 달성합니다.

이렇게 작은 모델(0.5B/1.5B)에서 얻은 강력한 결과는 올바른 기반이 매개변수 개수보다 중요하다는 저희의 논제를 입증합니다. 생성 모델은 코드 의미론을 이해하며 이러한 이해는 표현 작업으로 직접 이전됩니다.

이는 向量모델과 생성이 동일한 기반에서 나타나는 통합 아키텍처를 통해 검색 기반 모델로 가능한 것의 경계를 넓히는 Jina AI의 더 넓은 비전과 일치합니다.