拡散モデル系のjev(AIメモ)
github.com/looskis/tarski
https://x.com/thsottiaux/status/2104986448269279399
GitHub - looskis/tarski: Reading diffusi
tarskiは、拡散言語モデル(DiffusionGemmaやLLaDAなど)を用いた意思決定(分類やスコアリングなど)の「読み取りポリシー(Read Policy)」を最適化するためのCLI(コマンドライン)ツールです。
モデルの重みを一切変更(追加学習)することなく、プロンプトの工夫や設定の最適化だけで、推論の精度向上と高速化(OpenJev比で約1.8倍速、パス数を3.65から1に削減)を実現するための機能を提供します。
主な機能と特徴は以下の通りです:
最適な読み取り設定の自動探索 (tarski tune): ユーザーが用意した少量の校正データ(JSONL形式)を元に、「プロンプトのレイアウト(例:判定対象テキストを先に配置する)」「設問の並び順」「解答スロットの初期化方法」の最適な組み合わせを自動で選択します。
予測信頼度の補正: 質問タイプごとに負の対数尤度(NLL)を計算し、温度パラメータ(Temperature)を調整することで、モデルの出力する確率の歪みを補正(キャリブレーション)します。
軽量なオフライン計算: 最適な組み合わせを選ぶ tune や、精度を検証する eval といったコマンドは、モデル本体をメモリにロードせず、Python標準ライブラリのみを用いて高速にオフライン処理を実行します。
ユーザー自身のデータへの適用 : チケット分類(例:請求、サポート)や緊急度判定などの任意のJSONデータ構造に対して、設定ファイル(JSON)を介して簡単にこの最適化フローを適用し、最終的な予測(tarski predict)を出力させることができます。
つまり、拡散言語モデルの振る舞いに関する最新の研究成果を、ユーザー自身のタスク(カスタマーサポートの自動分類など)に手軽かつ低コストで導入・最適化するための実用フレームワークとして機能します。
背景として、
2026年9月にTypeSafe AI社から「Jev」という、自然言語による分類やスコアリング、構造化データの出力などに特化したモデルが公開されました。その後、開発者のrazorback16氏らによって、DiffusionGemmaなどの拡散LLMを利用してJev互換の処理を行う「OpenJev」というプロジェクトが登場しました。
note
いまご覧になっているリポジトリ「tarski」は、そのOpenJevモデル(DiffusionGemmaなど)の重みは一切変更せずに、読み取りポリシーを最適化するツールです。OpenJevのデフォルト実装では出力決定までに複数回(平均3.65回)のパスが必要だったところを、このtarskiを導入することで1パス(約1.8倍高速)で処理できるようにし、さらに精度も向上させるアプローチをとっています。
後学習は無しで強いの?
モデルの重みに対する追加学習(ファインチューニングやLoRAなど)は一切行っていません。モデルの重みは完全に固定(frozen)されたままです。
学習を行わずに精度が向上し(+5.3ポイント)、かつ高速化(約1.8倍)できている理由は、拡散言語モデル特有の性質に合わせた「推論時ポリシー(Read Policy)」の最適化に特化しているためです。
精度向上と1パス化を実現している主な工夫は以下の通りです。
-
解答スロットの初期化方法(Slot Initialization)
拡散モデルは初期ノイズからトークンを復元していきますが、OpenJev標準のランダムトークン初期化に対し、tarskiでは語彙埋め込みの平均値(vocabulary-mean embedding)を初期スロットに配置します。これにより初期状態のノイズの偏りが抑えられ、DiffusionGemmaにおいて復元の安定性が大幅に向上します。 -
プロンプトの構造配置(Prompt Layout)
拡散モデルは自己回帰型(Autoregressive)モデルと異なり、双方向にアテンションが働くため入力の配置順に敏感です。判定対象となるテキストを先頭に配置するstate_firstレイアウトを採用することで、後続の質問に対する順序依存性や干渉を大幅に緩和しています。 -
設問の並び順(Question Order)の最適化
一度に複数の判定を行う際、質問をどの順序でキャンバスに並べるか(元の順序、逆順、ローテーションなど)を少量の校正データ上で探索し、モデルが最も相互干渉を起こしにくい順序を選択します。 -
信頼度温度の校正(Confidence Calibration)
質問のタイプごとに負の対数尤度(NLL)を用いて温度パラメータ(Temperature)を最適化し、予測確率の歪みを補正(キャリブレーション)しています。
なぜ「3.65パス → 1パス」に減らせるのか
OpenJevのデフォルト実装では、確信度が低い場合や不整合が起きた際に自動で再読み取り(re-read)を行うため、平均して3.65回のデコーダーパス(約995ms)を要していました。
tarskiでは上記の入力設計によって初回の1パス目で十分な精度(65.9% → 71.2%)が出るようになるため、再読み取りを行わず1パス(約560ms)で処理を打ち切ることが可能になっています。
なお、最適な設定を探索する tarski tune 自体も、モデルを動かすのではなく収集済みの読み取り結果(JSONL)に対してPython標準ライブラリだけで最適な組み合わせを計算するため、手元で手軽に実行できる設計になっています。

tarskiが拡散言語モデル(DLM、ここではDiffusionGemmaなど)を用いて、追加学習(ファインチューニング)なしで高精度・高速な推論を実現する仕組みをまとめたインフォグラフィックです。大きく3つのセクションに分かれています。
tarskiによる計算プロセス構造図
Section 1: 現状の課題(OpenJev標準実装)
従来のOpenJev実装では、標準的なプロンプト、ランダムなスロット、そして確信度が低い場合の自動再読み取り(平均3.65パス)によって推論を行っており、結果として低精度で低速でした。
Section 2: tarskiによるオフラインチューニング(学習なし)
ユーザーの手元で、少量の校正データ(JSONL)を用いて最適な設定を探索します。
tarski collectコマンドで、様々なプロンプト形式、質問順、スロットを試し、raw確率を収集します。
tarski tuneコマンドで、収集した確率に基づき、精度を最大化する設定(fitted JSON)を選択し、質問タイプごとの温度(Confidence Temperature)を校正(キャリブレーション)します。このプロセス全体は、モデルをメモリにロードせずに行われます。
Section 3: 最適化された推論プロセス(tarski predict)
新しいデータに対して、チューニングで決定した最適なポリシーを適用します。
「解答スロットの初期化(語彙平均値)」「プロンプト構造(state_first)」「質問の並び順」を最適化した状態で、 frozen(固定)された拡散モデルによって高速な1パス推論を行います。
最後に、保存された温度を適用して、校正された予測分布を出力します。
図の右側にある比較表(Passes, Accuracy, Latency, Status)が示す通り、tarskiを導入することで、モデルの重みを一切変更せずに、パス数を3.65から1に削減し、精度を5.3ポイント向上、推論を約1.8倍高速化できることが視覚的にわかります。
このように、tarskiは拡散言語モデル独自の性質を理解し、推論時の「読み取り方」を最適化することで、学習なしで高いパフォーマンスを引き出すためのツールです。
Write a comment