Edit model card
YAML Metadata Warning: empty or missing yaml metadata in repo card (https://huggingface.co/docs/hub/model-cards#model-card-metadata)


Difference-based Contrastive Learning for Korean Sentence Embeddings

Quick tour

import torch
from transformers import AutoModel, AutoTokenizer

def cal_score(a, b):
    if len(a.shape) == 1: a = a.unsqueeze(0)
    if len(b.shape) == 1: b = b.unsqueeze(0)

    a_norm = a / a.norm(dim=1)[:, None]
    b_norm = b / b.norm(dim=1)[:, None]
    return torch.mm(a_norm, b_norm.transpose(0, 1)) * 100

model = AutoModel.from_pretrained('BM-K/KoDiffCSE-RoBERTa')
tokenizer = AutoTokenizer.from_pretrained('BM-K/KoDiffCSE-RoBERTa')

sentences = ['์น˜ํƒ€๊ฐ€ ๋“คํŒ์„ ๊ฐ€๋กœ ์งˆ๋Ÿฌ ๋จน์ด๋ฅผ ์ซ“๋Š”๋‹ค.',
             '์น˜ํƒ€ ํ•œ ๋งˆ๋ฆฌ๊ฐ€ ๋จน์ด ๋’ค์—์„œ ๋‹ฌ๋ฆฌ๊ณ  ์žˆ๋‹ค.',
             '์›์ˆญ์ด ํ•œ ๋งˆ๋ฆฌ๊ฐ€ ๋“œ๋Ÿผ์„ ์—ฐ์ฃผํ•œ๋‹ค.']

inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")
embeddings, _ = model(**inputs, return_dict=False)

score01 = cal_score(embeddings[0][0], embeddings[1][0])  # 84.56
# '์น˜ํƒ€๊ฐ€ ๋“คํŒ์„ ๊ฐ€๋กœ ์งˆ๋Ÿฌ ๋จน์ด๋ฅผ ์ซ“๋Š”๋‹ค.' @ '์น˜ํƒ€ ํ•œ ๋งˆ๋ฆฌ๊ฐ€ ๋จน์ด ๋’ค์—์„œ ๋‹ฌ๋ฆฌ๊ณ  ์žˆ๋‹ค.'
score02 = cal_score(embeddings[0][0], embeddings[2][0])  # 48.06
# '์น˜ํƒ€๊ฐ€ ๋“คํŒ์„ ๊ฐ€๋กœ ์งˆ๋Ÿฌ ๋จน์ด๋ฅผ ์ซ“๋Š”๋‹ค.' @ '์›์ˆญ์ด ํ•œ ๋งˆ๋ฆฌ๊ฐ€ ๋“œ๋Ÿผ์„ ์—ฐ์ฃผํ•œ๋‹ค.'


Python Pytorch

Encoder Models

Baseline encoders used for korean sentence embedding - KLUE-PLMs

Model Embedding size Hidden size # Layers # Heads
KLUE-BERT-base 768 768 12 12
KLUE-RoBERTa-base 768 768 12 12

Large pre-trained models need a lot of GPU memory to train


The data must exist in the "--path_to_data" folder

Training - unsupervised

python main.py \
    --model klue/roberta-base \
    --generator_name klue/roberta-small \
    --multi_gpu True \
    --train True \
    --test False \
    --max_len 64 \
    --batch_size 256 \
    --epochs 1 \
    --eval_steps 125 \
    --lr 0.00005 \
    --masking_ratio 0.15 \
    --lambda_weight 0.005 \
    --warmup_ratio 0.05 \
    --temperature 0.05 \
    --path_to_data Dataset/ \
    --train_data wiki_corpus_examples.txt \
    --valid_data valid_sts.tsv \
    --ckpt best_checkpoint.pt
bash run_diff.sh

Using roberta as an encoder is beneficial for training because the KoBERT model cannot build a small-sized generator.


python main.py \
    --model klue/roberta-base \
    --generator klue/roberta-small \
    --train False \
    --test True \
    --max_len 64 \
    --batch_size 256 \
    --path_to_data Dataset/ \
    --test_data test_sts.tsv \
    --path_to_saved_model output/best_checkpoint.pt

Performance - unsupervised

Model Average Cosine Pearson Cosine Spearman Euclidean Pearson Euclidean Spearman Manhattan Pearson Manhattan Spearman Dot Pearson Dot Spearman
KoSRoBERTa-baseโ€  N/A N/A 48.96 N/A N/A N/A N/A N/A N/A
KoSRoBERTa-largeโ€  N/A N/A 51.35 N/A N/A N/A N/A N/A N/A
KoSimCSE-BERT 74.08 74.92 73.98 74.15 74.22 74.07 74.07 74.15 73.14
KoSimCSE-RoBERTa 75.27 75.93 75.00 75.28 75.01 75.17 74.83 75.95 75.01
KoDiffCSE-RoBERTa 77.17 77.73 76.96 77.21 76.89 77.11 76.81 77.74 76.97


This work is licensed under a Creative Commons Attribution-ShareAlike 4.0 International License.

Creative Commons License


   title={{DiffCSE}: Difference-based Contrastive Learning for Sentence Embeddings},
   author={Chuang, Yung-Sung and Dangovski, Rumen and Luo, Hongyin and Zhang, Yang and Chang, Shiyu and Soljacic, Marin and Li, Shang-Wen and Yih, Wen-tau and Kim, Yoon and Glass, James},
   booktitle={Annual Conference of the North American Chapter of the Association for Computational Linguistics (NAACL)},
      title={KLUE: Korean Language Understanding Evaluation},
      author={Sungjoon Park and Jihyung Moon and Sungdong Kim and Won Ik Cho and Jiyoon Han and Jangwon Park and Chisung Song and Junseong Kim and Yongsook Song and Taehwan Oh and Joohong Lee and Juhyun Oh and Sungwon Lyu and Younghoon Jeong and Inkwon Lee and Sangwoo Seo and Dongjun Lee and Hyunwoo Kim and Myeonghwa Lee and Seongbo Jang and Seungwon Do and Sunkyoung Kim and Kyungtae Lim and Jongwon Lee and Kyumin Park and Jamin Shin and Seonghyun Kim and Lucy Park and Alice Oh and Jungwoo Ha and Kyunghyun Cho},
  title={KorNLI and KorSTS: New Benchmark Datasets for Korean Natural Language Understanding},
  author={Ham, Jiyeon and Choe, Yo Joong and Park, Kyubyong and Choi, Ilji and Soh, Hyungjoon},
  journal={arXiv preprint arXiv:2004.03289},
Downloads last month
Model size
111M params
Tensor type