🚀 ModularStarEncoder-1B 事前学習モデル
ModularStarEncoder-1Bは、コードエンコーディングに特化した事前学習モデルです。5つの出口点を持つモジュール型構造を持ち、下流タスクに応じた柔軟な微調整が可能です。
🚀 クイックスタート
ModularStarEncoder-1Bは、The Stack v2 で事前学習されたエンコーダです。ModularStarEncoderは5つの出口点を持つモジュール型の事前学習エンコーダで、下流タスクに応じて複数の出口微調整を行うことができます。
このモデルは StarCoder-2 をベースに構築され、bfloat16でパラメータ数を15Bから1Bに削減しています。アーキテクチャは36の隠れ層から構成され、各層には16のアテンションヘッドと4つのキーバリューヘッドがあり、Grouped Query Attention (GQA) を使用しています。
モデルは Rotary Positional Encoding (RoPE) を使用し、基本周期 theta = 10^-6 で、隠れ次元数は1024、中間サイズは12,288です。
効率を向上させるために、因果自己アテンション層を双方向自己アテンションに置き換えました。StarCoder-2 がスライディングウィンドウアテンションを使用するのとは異なり、完全アテンションを選択してより高いモジュール性を確保し、スライディングウィンドウメカニズムの受容野制約を回避しました。さらに、最大入力長を2048トークンに拡張し、以前のコードエンコーダ(StarEncoder など)よりも長いコードスニペットを扱えるようにしました。
最後に、推論速度を向上させるためにFlashAttention V2を統合しています。
コード間変換やテキストからコードへの微調整バージョンについては、ModularStarEncoder-finetuned を参照してください。
💻 使用例
基本的な使用法
from transformers import AutoModel
from transformers import AutoTokenizer
model = AutoModel.from_pretrained("andreagurioli1995/ModularStarEncoder", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("andreagurioli1995/ModularStarEncoder")
code_snippet = "your code to embed here"
sentence = f"{tokenizer.sep_token}{code_snippet}{tokenizer.cls_token}"
tokenized_sensence = tokenizer(sentence, return_tensors="pt",truncation=True, max_length=2048)
embedded_sentence = model(**tokenized_sensence)
出力として6つの要素が得られます。
- last_hidden_state: モデルの最後の隠れ層の表現
- hidden_states: モデルのすべての隠れ層からの生の表現(プーリング、正規化、投影なし)
- loss: 正解が与えられた場合の損失値(推論時はNone)
- prediction_logits: マスク言語モデリングヘッドからの予測スコア
- seq_relationship_scores: 文脈内損失の予測スコア(意味のあるスコアを得るには、区切りトークンで複数のサンプルを連結します)
- attentions: エンコーダからのアテンションスコア
🔧 技術詳細
学習
ModularStarEncoderは、バッチサイズ3.99Mトークンで245,000ステップの事前学習を行い、1Tトークンを処理しました。事前学習と微調整は、Leonardo スーパーコンピュータを使用して512台のNVIDIA Ampere (64GB) GPUで行われ、450,000 GPU作業時間が必要でした。
ハイパーパラメータ |
値 |
隠れ層のサイズ |
1024 |
最大位置埋め込み |
2048 |
アテンションヘッドの数 |
12 |
キーバリューヘッドの数 |
4 |
隠れ層の数 |
36 |
アテンション |
GQA |
パラメータ数 |
≈1B |
学習トークン数 |
≈1T |
損失関数 |
MLM + 文脈内損失 |
多層損失 |
あり |
📄 ライセンス
このモデルはBigCode OpenRAIL-M v1ライセンス契約の下で提供されています。完全な契約は こちら で確認できます。
📚 引用
@article{gurioli2025modeltrainallhierarchical,
title={One Model to Train them All: Hierarchical Self-Distillation for Enhanced Early Layer Embeddings},
author={Andrea Gurioli and Federico Pennino and João Monteiro and Maurizio Gabbrielli},
year={2025},
eprint={2503.03008},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2503.03008},
}