ByteDance Seed、連続潜在LM「Cola2」を公開　8タスク平均60.10

ByteDance-Seedは、連続潜在表現と離散テキスト生成を組み合わせる言語モデル「Cola2」をHugging FaceでByteDance-Seed/Cola2として公開しました。公開チェックポイントは346.03Bソーストークンで学習しており、8タスク平均60.10を報告しています。

Hugging Faceは、モデルの重みやデータセットを共有するHugging Face Hubを提供する企業です。

何が新しいか

Cola2は、テキストを意味的なスパンに動的に分割し、各スパンを圧縮された連続潜在表現に変換します。block-causal Diffusion Transformer(DiT)が潜在ブロックを生成し、span-causal token decoderが離散トークンへ変換します。構成要素はdynamic-compression VAE、block-causal DiT、潜在フィードバック付きspan-causal token decoderの3つです。

学習は2段階です。Stage Iでdynamic-compression encoderとtoken decoderを同時に学習します。Stage IIではトークン側の教師信号でVAEの適応を続けつつ、約4BのDiTがFlow Matchingで切り離された潜在分布を学習します。推論では潜在ブロックの提案、トークン復号、確定テキストの固定スパン配置下での再エンコードを繰り返し、連続履歴と離散履歴を整合させます。訓練コードは公開物に含まれません。

ベンチマーク

各値は潜在ノイズを変えた3回の評価の平均で、単位は正解率(%)です。学習トークン行のみトークン数です。

タスク	Llama-3-8B	Qwen2.5-7B	Qwen3-8B-Base	LLaDA-8B-Base	Cola2
事前学習トークン	15T+	18T	~36T	2.3T	346B
RACE	75.77	89.95	90.65	87.40	69.91
NarrativeQA	42.70	32.90	37.30	36.00	32.70
HellaSwag	39.21	75.67	68.13	69.65	41.75
StoryCloze	79.49	95.09	97.44	95.94	80.42
LAMBADA	75.70	71.82	71.20	67.90	63.00
ARC-Easy	90.86	96.00	97.35	93.85	87.62
TruthfulQA	63.04	69.03	72.58	64.87	54.10
MMLU	63.19	70.16	72.31	61.27	51.32
8タスク平均	66.25	75.08	75.87	72.11	60.10

ByteDance-Seedは、潜在DiTが未収束の状態でHellaSwagとStoryClozeでLlama-3-8Bを上回り、報告上の事前学習トークンは43分の1以下だとしています。一方で、各モデルは学習レシピ、トークナイザー、トークン数の数え方が異なり、この表は生成方式の差を切り分けるものではなく能力の目安だと説明しています。数値は論文の結果で、潜在生成が確率的なため個別の実行では結果が変わる可能性があります。

モデルサイズ・データ・評価設定をそろえた比較(AR、LLaDA、Cola DLM、Cola-MDM、Cola2-MDM)では、Cola-MDMのスケーリングが最も強く、Cola2-MDMは収束が遅いものの約150BトークンからLLaDAとの差を縮め続け、約210Bトークン時点でも飽和していません。decoderのclassifier-free guidance最適重みは学習の大半で-0.9付近にあり、徐々に0へ近づきます。

要件と試し方

• Python 3.10+、PyTorch 2.2+、Transformers 4.51–4.x、safetensors 0.4+
• 推奨環境はCUDAとbfloat16
• カスタムTransformersコードのためtrust_remote_code=Trueが必須

読み込みはAutoTokenizer.from_pretrained(model_id, trust_remote_code=True)とAutoModel.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True)で行い、.to("cuda").eval()でCUDAへ移したうえでmodel.generate(**inputs, generator=generator)を実行します。サンプルではtorch.Generatorにmanual_seed(42)を与えています。

Hugging Faceのリポジトリには8分割のsafetensors重み、トークナイザーと生成設定、remote-codeモジュール、8評価タスク用の正規化JSONLデータ、SHA256SUMSが含まれます。推論実装、コマンドラインツール、サンプル、ベンチマークランナー、文書はCola2のGitHubリポジトリで提供されます。