ECeLLM: Training a Custom E-Commerce Embedding Model

E-Commerce LLM Embedding Model Training Log (prj_ecellm)

Overview

This project covers training a custom embedding model specialized for the e-commerce domain. Over 42 commits from December 19, 2024 to January 23, 2025, it implements a SentenceTransformer-based custom embedding model and a CrossEncoder reranker.

Motivation

General-purpose embedding models (OpenAI, BGE, etc.) work well on general text, but fall short on e-commerce product data:

• "아이폰 15 프로 맥스 256GB 블루" vs "Apple iPhone 15 Pro Max" — same product, different surface forms
• Product attributes (brand, capacity, color) carry different weight than they do in general text
• Korean product names follow patterns that general models don't handle well

Development

Phase 1: Foundation (12/19–12/26)

Core infrastructure was built in the first week after project initialization.

# 커스텀 토크나이저 설정
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("klue/bert-base")
# 이커머스 특화 토큰 추가
tokenizer.add_tokens(["[BRAND]", "[CATEGORY]", "[SPEC]"])

Commit messages during this period were heavily focused on tokenizer work: Tokenizer_tool, del_local_tokenizer, del_token, and similar.

Phase 2: SentenceTransformer Training (12/27–01/06)

A custom SentenceTransformer was implemented on top of BERT-NLI.

from sentence_transformers import SentenceTransformer, losses

# 커스텀 SentenceTransformer 모델
model = SentenceTransformer("klue/bert-base")

# Contrastive Loss로 학습
train_loss = losses.CosineSimilarityLoss(model)

model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=10,
    warmup_steps=100,
    evaluator=evaluator
)

The bert_nli_add_custom_sentencetransformer commit on January 6 completed custom training with NLI data.

Phase 3: StableAdamW & Optimization (01/07–01/09)

StableAdamW was adopted to improve training stability.

from torch.optim import AdamW  # 기존
# ↓
from custom_optimizers import StableAdamW  # 적용

optimizer = StableAdamW(
    model.parameters(),
    lr=2e-5,
    weight_decay=0.01,
    eps=1e-8
)

Final stabilization landed in the 0109_fix_StableAdamW commit.

Phase 4: CrossEncoder Reranker (01/12–01/14)

A two-stage pipeline was implemented: a Bi-Encoder (SentenceTransformer) retrieves candidates, then a CrossEncoder performs precise reranking.

from sentence_transformers import CrossEncoder

# CrossEncoder 리랭커
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

# 2단계 검색 파이프라인
candidates = bi_encoder.encode(query, top_k=100)  # 1차: 빠른 검색
reranked = reranker.predict([(query, doc) for doc in candidates])  # 2차: 정밀 리랭킹

Phase 5: Multi-Task Training (01/14–01/23)

MLM (Masked Language Modeling) and STS (Semantic Textual Similarity) tasks were added to give the model broader language understanding.

# MLM 태스크 추가 (01/15: mlm_task_add)
mlm_trainer = MLMTrainer(model, tokenizer, mlm_probability=0.15)

# STS 태스크 추가 (01/14: add_sts_trainer)
sts_trainer = STSTrainer(model, sts_dataset, loss=CosineSimilarityLoss)

# CE(Cross-Entropy) 트레이너 (01/14: st_ce_trainer)
ce_trainer = CETrainer(model, ce_dataset, loss=CrossEntropyLoss)

End-to-End Training Pipeline

상품 데이터
    │
    ├──▶ MLM 학습 ──▶ 도메인 지식 습득
    │
    ├──▶ STS 학습 ──▶ 유사도 이해
    │
    ├──▶ NLI 학습 ──▶ 의미 관계 이해
    │
    └──▶ CE 학습  ──▶ 분류 정확도
         │
         ▼
    커스텀 E-Commerce 임베딩 모델
         │
         ├──▶ Bi-Encoder: 빠른 시맨틱 검색
         └──▶ CrossEncoder: 정밀 리랭킹

Tech Stack

• PyTorch + Transformers (HuggingFace)
• SentenceTransformers (Bi-Encoder/CrossEncoder)
• KLUE BERT (Korean base model)
• StableAdamW (training stability)
• DeepSpeed (distributed training)

Results & Retrospective

Across 42 commits, the project achieved meaningful performance gains over general-purpose models on e-commerce domain tasks. The Bi-Encoder + CrossEncoder two-stage pipeline in particular struck a good balance between retrieval accuracy and latency. This work fed directly into the XGen platform's RAG pipeline and the SentenceTransformer training functionality in POLAR Trainer.