ko-sroberta-sts-mrl

jhgan/ko-sroberta-sts-mrl

This is a sentence-transformers model fine-tuned with Matryoshka Representation Learning (MRL) on top of klue/roberta-base. The model produces 768-dim embeddings, but the first m dims for m ∈ {768, 512, 256, 128, 64, 32} are themselves valid sentence representations — you can slice the embedding to trade accuracy for storage/latency without retraining.

• Training data: KorSTS
• Base model: klue/roberta-base
• Nested dims: 768, 512, 256, 128, 64, 32

KorSTS test set results

All values are reported as percentages (×100).

Usage

Full 768-dim embedding

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("jhgan/ko-sroberta-sts-mrl")
embeddings = model.encode(["안녕하세요", "반갑습니다"])
print(embeddings.shape)  # (2, 768)

Truncated embedding (recommended pattern)

import torch.nn.functional as F
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("jhgan/ko-sroberta-sts-mrl")
emb = model.encode(["안녕하세요", "반갑습니다"], convert_to_tensor=True)

# Slice to the first 64 dims and re-normalise for cosine similarity
emb_64 = F.normalize(emb[:, :64], p=2, dim=1)

Or fix the truncation at load time (requires sentence-transformers >= 2.7.0):

model = SentenceTransformer("jhgan/ko-sroberta-sts-mrl", truncate_dim=64)

Citation

@inproceedings{kusupati2022matryoshka,
    title     = {Matryoshka Representation Learning},
    author    = {Kusupati, Aditya and Bhatt, Gantavya and Rege, Aniket and
                 Wallingford, Matthew and Sinha, Aditya and Ramanujan, Vivek and
                 Howard-Snyder, William and Chen, Kaifeng and Kakade, Sham and
                 Jain, Prateek and Farhadi, Ali},
    booktitle = {Advances in Neural Information Processing Systems},
    year      = {2022},
    url       = {https://arxiv.org/abs/2205.13147}
}

This model is part of the ko-sentence-transformers project; see the repository for training scripts and the non-MRL baselines (jhgan/ko-sroberta-sts, jhgan/ko-sroberta-nli, jhgan/ko-sroberta-multitask).