kurogane/ca-reward-distill-1B-ja-preview

ca-reward-distill-1B-ja

本コード、モデルはすべてChatGPTのみを使用して作成されました。よって、全内容を公開します。
同様に、動作確認や品質保証はできていません。使用する場合は丁寧に動作確認してください。

ca-reward-distill-1B-ja は、日本語の prompt-response ペアに対して単一のスカラー報酬スコア を出力する 1B クラスの Reward Model です。ベースモデルには sbintuitions/sarashina2.2-1b-instruct-v0.1 を使い、教師モデル cyberagent/ca-reward-3b-ja が付与したスコアを回帰蒸留する形で学習しました。

このモデルは 生成モデルではありません。用途は、応答候補の順位付け、ベスト応答の選択、データフィルタリング、蒸留用の補助スコアリングなどを想定しています。

Model Details

Summary

• Task: 日本語 Reward Modeling / response scoring
• Architecture: AutoModelForSequenceClassification (num_labels=1, regression)
• Base model: sbintuitions/sarashina2.2-1b-instruct-v0.1
• Teacher model: cyberagent/ca-reward-3b-ja
• Training method: score regression distillation
• Primary language: Japanese

What the score means

出力は「この応答が、この prompt に対してどれだけ好ましいか」を表す相対的なスコアです。絶対値そのものよりも、同一 prompt に対する複数候補の相対比較で使うことを推奨します。

このリポジトリの学習・推論コードでは、以下の 2 種類の値を扱います。

• student_score_norm: 学習時の正規化空間での出力
• student_score_denorm: score_normalization.json を使って教師スコア空間へ戻した値

多くの用途では、同一 prompt 内で大小比較できれば十分です。しきい値判定よりも、順位付け用途に向いています。

Training Pipeline

このチェックポイントは、公開しているスクリプト群のうち 回帰蒸留の本線で作成したものです。

1. build_mixed_prompt_dataset.py で、日本語中心の prompt-only データセットを構築
2. generate_teacher_dataset.py で、各 prompt に対して複数候補応答を生成し、cyberagent/ca-reward-3b-ja で採点
3. train_student_rm_regression.py で、教師スコアを回帰ラベルとして生徒 RM を学習
4. evaluate_student_rm_against_teacher.py で、生徒 RM と教師スコアの整合性を評価

今回公開しているチェックポイントは pairwise 学習ではなく、score regression distillation で学習したモデルです。

Training Data

Prompt pool

公開している build summary に基づく prompt pool は 51,300 prompt です。内訳は次の通りです。

Source	Rows
llm-jp/magpie-sft-v1.0	15,000
llm-jp/extraction-wiki-ja (v0.3)	10,000
llm-jp/wizardlm8x22b-logical-math-coding-sft-ja	8,000
llm-jp/oasst2-33k-ja	6,000
llm-jp/oasst1-21k-ja	3,000
llm-jp/Synthetic-JP-EN-Coding-Dataset (Japanese prompts only)	3,000
llm-jp/databricks-dolly-15k-ja	2,000
llm-jp/llm-jp-instructions (v1.0/train)	300
llm-jp/oasst2-33k-en → Japanese-answer wrapped	2,000
llm-jp/oasst1-21k-en → Japanese-answer wrapped	2,000

英語プロンプトを使う一部ソースでは、プロンプト末尾に “Please answer in Japanese. / 日本語で回答してください。” を追加し、日本語応答を促す形にしています。

Teacher-scored dataset

公開している実行コマンドでは、生成モデルとして sbintuitions/sarashina2.2-1b-instruct-v0.1 を使い、各 prompt につき 4 candidate responses を生成しています。そのため、教師データセットは 約 205,200 件 の prompt-response-score 行から構成されます。

各候補応答は cyberagent/ca-reward-3b-ja により採点され、teacher_score として保存されます。

Training Procedure

公開している実行コマンドでは、以下の設定で学習しています。

• max_length=2048
• per_device_train_batch_size=2
• gradient_accumulation_steps=16
• learning_rate=1e-5
• num_train_epochs=1
• gradient_checkpointing=True

加えて、スクリプト既定値として次が使われます。

• base model: sbintuitions/sarashina2.2-1b-instruct-v0.1
• score normalization: train_zscore
• loss: mse
• validation split: 入力に validation が無い場合、prompt_hash ベースの安定ハッシュ分割で validation_ratio=0.02

Training result

Convert to safetensors

以下のコマンドでsafetensorsに変換可能です。アップロードしたファイルは変換済みのものです。

python /mnt/data/convert_student_rm_to_safetensors.py \
  --model-dir ./student_rm_regression_trial/final_model \
  --output-dir ./student_rm_regression_trial/final_model_safetensors

Evaluation

このモデルは evaluate_student_rm_against_teacher.py により、教師 RM に対する近似精度として評価できます。主な評価指標は次の通りです。

• global Pearson / Spearman
• MSE / MAE
• prompt 単位の top1 agreement
• pairwise accuracy (micro / macro)
• prompt-wise Spearman

Metrics for this release

Metric	Value
Pearson (denorm)	0.9705
Spearman (denorm)	0.9631
MAE (denorm)	0.0881
MSE (denorm)	0.1929
Top-1 agreement	0.6833
Pairwise accuracy (micro)	0.8291
Pairwise accuracy (macro)	0.8282
Prompt-wise Spearman mean	0.7172

教師RMが付与したスコア付きデータセットの train split 上で評価したところ、本モデルは Pearson 0.9705、Spearman 0.9631 を達成した。 また、候補ランキングに関しては Top-1 agreement 0.6833、pairwise accuracy 0.8291 を示し、1B student model が教師 Reward Model を高精度に蒸留できていることを確認した。 ただし本評価は教師生成データ上での teacher-agreement を測るものであり、人手選好との一致を直接測るものではない。

評価結果の全量 JSON

{ "model_dir": "/media/kurogane/datasets/ ca-reward-distill/student_rm_regression_trial/final_model", "teacher_dataset_dir": "/media/kurogane/datasets/ ca-reward-distill/teacher_data_trial/final_dataset", "split": "train", "num_processes": 1, "group_column": "global_prompt_index", "score_column": "teacher_score", "score_normalization": { "clip_abs": null, "mean": 3.3561677871916498, "mode": "train_zscore", "std": 1.2300509459930353 }, "global_metrics": { "n": 205200, "teacher_score_mean": 3.3557875229182996, "teacher_score_std": 1.2309303641496014, "student_score_denorm_mean": 3.348719506911004, "student_score_denorm_std": 1.1915005575299094, "student_score_norm_mean": -0.0060552616173412364, "student_score_norm_std": 0.9686595188689493, "mse_denorm": 0.0881010115453273, "mae_denorm": 0.19286510218406958, "pearson_denorm": 0.9705123593037001, "spearman_denorm": 0.9631207945431493 }, "group_metrics": { "num_prompt_groups": 51300, "num_prompt_groups_multi_candidate": 51300, "candidate_count_per_prompt_mean": 4.0, "candidate_count_per_prompt_median": 4.0, "candidate_count_per_prompt_max": 4, "top1_agreement": 0.6832748538011696, "top1_agreement_num_prompts": 51300, "pairwise_accuracy_micro": 0.8291460903627568, "pairwise_accuracy_micro_num_pairs": 298685, "pairwise_accuracy_macro": 0.8281829008766277, "pairwise_accuracy_macro_num_prompts": 51219, "prompt_spearman_mean": 0.7171808839384306, "prompt_spearman_median": 0.7999999999999999, "prompt_spearman_num_prompts": 51300, "pairwise_margin": 0.0 } }

Important note on evaluation

この評価は、人手選好ラベルに対する直接評価ではなく、教師 RM に対する蒸留忠実度評価です。したがって、この README の指標は「教師をどれくらい再現できているか」を表します。

How to Use

Minimal example with transformers

import json
import torch
from pathlib import Path
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from huggingface_hub import hf_hub_download

model_id = "kurogane/ca-reward-distill-1B-ja"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id)
model.eval()
model.to(device)

prompt = "富士山について短く説明して"
response_a = "富士山は日本で最も高い山で、静岡県と山梨県にまたがる成層火山です。"
response_b = "富士山はたぶんどこかにある山です。"

def build_text(prompt: str, response: str) -> str:
    chat = [
        {"role": "user", "content": prompt},
        {"role": "assistant", "content": response},
    ]
    return tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=False)

texts = [build_text(prompt, response_a), build_text(prompt, response_b)]
inputs = tokenizer(
    texts,
    return_tensors="pt",
    truncation=True,
    max_length=2048,
    padding=True,
    add_special_tokens=False,
).to(device)

with torch.no_grad():
    scores_norm = model(**inputs).logits.squeeze(-1).float().cpu().tolist()

print("normalized scores:", scores_norm)

# Optional: denormalize to the teacher-score scale
try:
    score_stats_path = hf_hub_download(repo_id=model_id, filename="score_normalization.json")
    with open(score_stats_path, "r", encoding="utf-8") as f:
        score_stats = json.load(f)
    if score_stats.get("mode") == "train_zscore":
        mean = float(score_stats.get("mean", 0.0))
        std = float(score_stats.get("std", 1.0))
        scores_denorm = [s * std + mean for s in scores_norm]
        print("denormalized scores:", scores_denorm)
except Exception:
    pass

CLI example

公開している推論スクリプトを使う場合:

python score_student_rm_minimal.py \
  --model-dir ./final_model \
  --prompt "富士山について短く説明して" \
  --response "富士山は日本で最も高い山です。"

Intended Uses

• 同一 prompt に対する複数応答候補の順位付け
• best-of-n 生成の候補選択
• 応答品質の自動フィルタリング
• 蒸留・データ収集・自己改善ループにおける補助スコアラー

Out-of-Scope / Limitations

• 生成モデルではありません。単体で回答文を生成する用途には使えません。
• 教師 RM を近似するよう学習しているため、教師モデルのバイアスや限界を引き継ぐ可能性があります。
• 学習・評価の中心は日本語 prompt-response ですが、一部に英語プロンプト由来の wrapped data を含みます。
• しきい値ベースの絶対判定より、相対比較での利用を推奨します。
• 人手選好ベンチマークに対する性能は、この README ではまだ報告していません。
• 安全性・公平性・有害性評価の専用チューニングは限定的です。高リスク用途では、人間のレビューや追加の安全対策と併用してください。

Repository / Reproducibility

学習・評価・最小推論に使ったスクリプトは GitHub で公開しています。

• Repository: kuroganegames/ca-reward-distill-1B-scripts
• Main scripts:
  • build_mixed_prompt_dataset.py
  • generate_teacher_dataset.py
  • train_student_rm_regression.py
  • evaluate_student_rm_against_teacher.py
  • score_student_rm_minimal.py

同リポジトリには、再読み込み確認用のテストスクリプトや、pairwise 蒸留用データセット構築スクリプトも含まれています。

License

apache-2.0