Jina Code Embeddings: 0.5B および 1.5B で SOTA のコード検索を実現

Efficient Code Embeddings from Code Generation Models

jina-code-embeddings is a novel code embedding model suite designed to retrieve code from natural language queries, perform technical question-answering, and identify semantically similar code snippets across programming languages. It makes innovative use of an autoregressive backbone pre-trained on both text and code, generating embeddings via last-token pooling. We outline the training recipe and demonstrate state-of-the-art performance despite the relatively small size of the models, validating this approach to code embedding model construction.

arXiv.orgDaria Kryvosheieva

本日、新しいコード向量模型のスイートであるjina-code-embeddingsをリリースします。サイズは0.5Bと1.5Bパラメータの2種類で、両方に対して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モデルは、20%小さいにもかかわらず、Qwen3-Embedding-0.6Bを5パーセントポイント上回り、1.5Bモデルは、voyage-code-3（79.23%）に匹敵し、gemini-embedding-001（77.38%）を上回っています。これらは両方とも、アーキテクチャが非公開のプロプライエタリモデルです。

Model	Parameters	Overall AVG	MTEB Code AVG
<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	Unknown*	79.23%	79.84%
gemini-embedding-001	Unknown*	77.38%	76.48%
jina-embeddings-v4	3.8B	74.11%	74.87%
Qwen3-Embedding-0.6B	600M	73.49%	74.69%

*Closed-source models with undisclosed architecture

両方のモデルは、非対称検索のために、クエリとドキュメントの両方の役割をサポートする、5つのタスク固有の指示プレフィックスでトレーニングされています。たとえば、クエリを埋め込むにはnl2code_queryを、ドキュメントを埋め込むにはnl2code_documentを使用できます。

Task	Use Case	Instruction Prefix
`nl2code`	"How to read CSV" → `pandas.read_csv()`	"Find the most relevant code snippet given the following query:\n"
`qa`	Technical Q&A retrieval	"Find the most relevant answer given the following question:\n"
`code2code`	Finding similar implementations	"Find an equivalent code snippet given the following code snippet:\n"
`code2nl`	Code to documentation	"Find the most relevant comment given the following code snippet:\n"
`code2completion`	Autocomplete scenarios	"Find the most relevant completion given the following start of code snippet:\n"

tagTraining Recipe

事前トレーニング済みのコード生成モデルを埋め込みバックボーンとして使用します。Qwen2.5-Coder-0.5Bおよび1.5Bを基盤として構築された当社のモデルは、次の機能を備えています。

Feature	jina-code-embeddings-0.5b	jina-code-embeddings-1.5b
Base Model	Qwen2.5-Coder-0.5B	Qwen2.5-Coder-1.5B
Embedding Dimensions	896	1536
Matryoshka Dimensions	64, 128, 256, 512, 896	128, 256, 512, 1024, 1536
Max Sequence Length	32,768 tokens	32,768 tokens
Pooling Strategy	Last-token pooling	Last-token pooling
Attention	FlashAttention2	FlashAttention2
Data Type	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パーセントポイントの利点により、jina-embeddings-v2、v3、およびv4で確立した平均プーリングの伝統を打ち破ることになりました。より多くの検索モデルがデコーダー専用大規模言語モデル (LLM) を基盤として構築されるにつれて、ラストトークンプーリングは自然な選択になります。平均プーリングは単方向注意メカニズムとうまく整合しません。平均プーリングは機能し、初期段階でより簡単にトレーニングできることが多いですが（おそらく凸最適化の状況による）、私たちの実験では、ラストトークンプーリングが達成するパフォーマンスの天井を下回って停滞することが一貫して示されています。

tagGetting Started

両方のモデルは、Search Foundation APIおよびsentence-transformers、transformers、llama.cppなどの一般的なフレームワークを介してシームレスに動作します。

tagVia API

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

tagVia `sentence-transformers`

from sentence_transformers import SentenceTransformer

# Load the model (choose 0.5b or 1.5b)
model = SentenceTransformer(
    "jinaai/jina-code-embeddings-1.5b",
    model_kwargs={"torch_dtype": "bfloat16"},
    tokenizer_kwargs={"padding_side": "left"}
)

# Natural language to code
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};"]

# Generate embeddings with task-specific prefixes
query_embeddings = model.encode(queries, prompt_name="nl2code_query")
document_embeddings = model.encode(documents, prompt_name="nl2code_document")

# Compute similarity
similarity = model.similarity(query_embeddings, document_embeddings)

tagVia `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'])

tagマトリョーシカ向量模型カットオフ

両方のモデルは、次元[64, 128, 256, 512, 896]に対してマトリョーシカ表現学習でトレーニングされており、再計算せずに向量模型を切り捨てることができます。

# 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 のより広範なビジョンと一致しています。向量模型と生成が同じ基盤から生まれる統一されたアーキテクチャであり、検索基盤モデルで可能なことの限界を押し広げています。