jina-v5-h256-lora-classification-universal

Jina v5 H256 — Universal Classification LoRA

A single LoRA adapter (405KB) that improves classification performance across all MTEB Classification tasks, built on top of the compressed gomyk/jina-v5-h256-distilled-conv embedding model.

Unlike per-task adapters, this is one universal adapter that enhances the embedding space for classification in general — no task-specific fine-tuning needed at inference time.

LoRA Specification

MTEB Classification Results

Average: 73.20% → 82.98% (+9.78%p)

Baseline: MTEB default evaluation (logistic regression on frozen embeddings) + LoRA: Same evaluation, but with universal LoRA adapter merged into embeddings

Training Method

Multi-Task Classification with shared LoRA backbone:

                    ┌─ Head_Amazon(256→2)   ─→ CE loss
                    ├─ Head_Banking(256→77)  ─→ CE loss
Input → [EuroBERT   ├─ Head_IMDB(256→2)     ─→ CE loss
         + LoRA]    ├─ Head_MTOP(256→11)    ─→ CE loss
         → pool  ──►├─ Head_Massive_I(256→60)─→ CE loss
                    ├─ Head_Massive_S(256→18)─→ CE loss
                    ├─ Head_Toxic(256→2)    ─→ CE loss
                    └─ Head_Tweet(256→3)    ─→ CE loss

All heads share the same LoRA backbone.
After training, heads are discarded — only the LoRA adapter is kept.
The LoRA-enhanced embeddings are universally better for classification.

Usage

Option 1: Merge LoRA into base model

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer
from huggingface_hub import hf_hub_download

# Load base model
model = AutoModel.from_pretrained(
    "gomyk/jina-v5-h256-distilled-conv", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(
    "gomyk/jina-v5-h256-distilled-conv", trust_remote_code=True)

# Download and merge LoRA
lora_path = hf_hub_download(
    "gomyk/jina-v5-h256-lora-classification-universal",
    "lora_adapter.pt")
lora_state = torch.load(lora_path, map_location="cpu", weights_only=True)

idx = 0
for name, module in model.named_modules():
    for target in ["q_proj", "k_proj", "v_proj", "o_proj"]:
        child = getattr(module, target, None)
        if child is not None and isinstance(child, nn.Linear):
            A = lora_state[f"lora_{idx}_A"]
            B = lora_state[f"lora_{idx}_B"]
            scaling = lora_state[f"lora_{idx}_scaling"].item()
            child.weight.data += (scaling * (A @ B)).T
            idx += 1

# Now use as normal embedding model
model.eval()
inputs = tokenizer("This movie was great!", return_tensors="pt",
                    truncation=True, max_length=128)
with torch.no_grad():
    outputs = model(**inputs)
    hidden = outputs.last_hidden_state
    mask = inputs["attention_mask"].unsqueeze(-1).float()
    embedding = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)

print(embedding.shape)  # [1, 256]

Option 2: Use with downstream classifier

# After merging LoRA (see above), add your own classifier
from sklearn.linear_model import LogisticRegression

# Encode your training data
train_embeddings = []  # encode your texts with the merged model
train_labels = [...]

clf = LogisticRegression(max_iter=1000)
clf.fit(train_embeddings, train_labels)

# Predict
test_embedding = ...  # encode test text
prediction = clf.predict(test_embedding)

File Structure

.
├── README.md              # This file
├── lora_adapter.pt        # LoRA A/B matrices (405KB)
├── meta.json              # Training metadata
└── task_info.json         # Per-task class counts

How LoRA Works

For each of the 24 attention projections (4 per layer x 6 layers):

  Original:  y = W @ x                    W: [256, 256] — frozen
  With LoRA: y = W @ x + (16/8) * (x @ A) @ B

  A: [256, 8]  — learned down-projection
  B: [8, 256]  — learned up-projection

  At merge time: W_new = W + 2.0 * (A @ B)^T  — zero overhead at inference

License

Apache 2.0