TL;DR
- 対象読者と課題: モデルパラメータの肥大化と推論時メモリ帯域(Memory Bandwidth)の制約に直面するMLインフラエンジニアおよびモデルアーキテクト向け。重み共有型再帰(Looped Transformer)により、重みVRAMを抑制したまま計算グラフ深度を拡張し、論理推論能力を選択的に引き上げる設計手法を提示する。
- タイムリーな背景: OpenAI「GPT-6 Astra」の推論プロセス隠蔽疑惑(The Information報道)とチーフサイエンティストJakub Pachockiの反論を契機に、2026年9月の最新研究「SMELT」をはじめとする、計算量適合(Compute-Matched)条件下での再帰的深度スケーリング則の研究・検証が加速している。
- 定量的効果と環境要件: Python 3.10+, PyTorch 2.4.0+(CUDA 12.4+)。最新のSMELT実験では従来型Transformer比で同等検証損失の達成に必要な訓練計算量を6.8〜18%削減し、テスト時のKLダイバージェンス判定による動的早期終了により難問への計算資源配分を最適化する。
背景と技術的課題
大規模言語モデル(LLM)のスケーリングにおいて、パラメータ数の単純な増加はメモリ容量と分散並列通信の物理的限界に直面している。従来のアーキテクチャでは、推論性能を高めるために層数(Depth)を増やすと重みパラメータが線形に増加し、GPUクラスタにおける重み格納VRAMとテンソル並列のインターコネクト帯域がボトルネック化する。さらに、推論時のテキストChain-of-Thought(CoT)によるテスト時計算量(Test-Time Compute)の拡張は、出力トークン数の増大に伴うレイテンシ悪化、KVキャッシュの肥大化、コンテキストウィンドウ枯渇を招く。
2026年9月上旬、米The Informationが報じたOpenAIの次世代モデル「GPT-6 Astra」に関するスクープは、MLエンジニアコミュニティで大きな議論を呼んだ。報道では、Astraが「再帰的深度(Recurrent Depth / Looped Transformer)」を採用し、同一のTransformerブロック群に隠れ状態をループ通過させることで、テキスト思考ログ(CoT)を外部に出力せず内部潜在空間で多段階推論を完結させ、推論プロセスを意図的に隠蔽しているのではないかというセキュリティ・監査性の懸念が提起された。
これに対し、OpenAIチーフサイエンティストのJakub Pachockiは声明を発表し、「Astraの計算グラフ深度はGPT-4の2倍以内にとどまる。推論トレースの短縮はモデル能力向上に伴う手戻り探索(Backtracking)の減少によるものであり、アーキテクチャ変更に起因する隠蔽ではない」と反論した。この論争の本質は、Transformerブロックの再帰的ループ適用が、モデルの「知識記憶容量」と「論理推論能力」にどのような非対称な影響を与えるのか、そして計算資源とメモリのトレードオフがどう規定されるかというスケーリング則の問題に帰着する。
コアアーキテクチャ・仕組みの解説
Looped Transformerの基本思想は、同一のTransformerブロック群(Attention、Feed-Forward Network、正規化層、残差接続)の重みを固定したまま、中間隠れ状態(Hidden States)を複数回再帰的に通過させる点にある。このアプローチの源流は2018年のUniversal Transformersに遡るが、2025年以降の研究によって実用的なスケーリング則が体系化された。
+-------------------------------------------------------------------------+
| Looped Transformer Dataflow & Computation |
+-------------------------------------------------------------------------+
Input Tokens: [ x_1, x_2, ..., x_t ]
│
▼
┌─────────────────────┐
│ Token Embedding & │
│ Initial Layers L_in │
└──────────┬──────────┘
│ Hidden State h_0
▼
┌───────────────────────────────────────────┐
│ Looped Transformer Stack (Weight-Shared) │ ◄─────────────┐
│ ┌─────────────────────────────────────┐ │ │
│ │ Layer 1 (Self-Attention + FFN/MoE) │ │ │
│ └──────────────────┬──────────────────┘ │ │
│ ▼ │ │
│ ┌─────────────────────────────────────┐ │ │ Loop Pass
│ │ Layer K (Self-Attention + FFN/MoE) │ │ │ (Pass 1..N)
│ └──────────────────┬──────────────────┘ │ │
└──────────────────────┼────────────────────┘ │
│ Hidden State h_k │
▼ │
[ Halting / Routing Gate ] │
│ - KL Divergence < eps? │
│ - Max passes reached? │
├──────────────────┬──────────────────────────┘
│ Yes (Exit) │ No (Next Pass)
▼ └────────► h_in = Proj(h_k || h_0)
┌─────────────────────┐
│ Final Layers L_out │
│ & LM Output Head │
└──────────┬──────────┘
▼
Next Token Logits / Prediction y_t
1. 重み共有と計算グラフのアンロール
オープンウェイトモデル「Nanbeige4.2-3B」(2026年7月公開)では、22層のTransformerブロックスタックを2周(2 passes)適用する。計算グラフを時間軸方向にアンロールすると44層相当の計算深度を持つが、保持する重みパラメータは22層分にとどまる。Nanbeigeの報告によれば、既存事前学習モデルのアップサイクリング(事後ループ化)よりもスクラッチからの事前学習が明確に優れており、3パス以上は計算コストに対する性能向上が鈍化し最適化が不安定化するため、2パスが最も費用対効果に優れる。
2. 知識記憶(Memorization)と論理推論(Reasoning)の機能分離
2025年6月の研究「Beyond Parameters: Exploring Virtual Logic Depth for Scaling Laws」(Zhu et al.)は、Looped Transformerの根本的な性質を実証した。実験の結果、ループ回数を増やして仮想論理深度(Virtual Logic Depth)を拡張しても、事実情報や知識の記憶容量はほぼ増加しないことが判明した。記憶容量は独立した実パラメータ数に厳密に束縛される。一方で、多段階の数学的推論やコード生成タスクでは、パラメータ数を固定したままループ回数を増やすだけで精度が大幅に向上する。すなわち、パラメータ数は「データベース(知識の保存)」として機能し、再帰的深度は「プロセッサ(論理演算サイクル)」として機能するという役割分担が成立する。
3. SMELT:計算量適合MoEスケーリング則
2026年9月発表の論文「SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers」(arXiv:2609.01343)は、計算量を公平に揃えた比較を実施した。中間半分のブロックを2回ループさせる構成に対し、ループによるFLOPs増大をモデル幅(隠れ次元)の縮小で相殺し、減少したパラメータ数をMoE(Mixture of Experts)のエキスパート増設で補う。さらにKVキャッシュサイズも従来型と同一に適合させた。最大54B非埋め込みパラメータ規模の検証において、SMELTは従来型Transformer比で6.8〜18%少ない訓練計算量(Compute Budget)で同等の検証損失を達成することを証明した。
4. 潜在推論(Latent Reasoning)と適応的早期終了
Geiping et al. (2025) の「Latent Reasoning」アーキテクチャでは、共有スタックの各ループ開始時に入力埋め込み表現を残差的に再結合し、テスト時の動的適応(Adaptive Halting)を導入した。連続する2つのループパスから出力される語彙確率分布のKLダイバージェンス(Kullback-Leibler Divergence)を計算し、確率分布の変化がしきい値未満となった時点で追加ループを打ち切る。平易なトークンでは1〜2ループで即時脱出し、複雑な推論を要するトークンでのみ最大ループまで演算を継続することで、推論FLOPsを自律的に配分する。
実践ハンズオン:動く実装コードとパイプライン
以下は、重み共有型スタック、入力射影残差、および推論時のKLダイバージェンスに基づく適応的早期終了(Adaptive Halting)を統合した最小構成のPyTorch実装である。外部ライブラリに依存せず、標準的なPyTorch環境で直接実行できる。
# 環境要件: Python 3.10+, torch 2.4.0+ (CUDA 12.4推奨)
import torch
import torch.nn as nn
import torch.nn.functional as F
class TransformerBlock(nn.Module):
"""標準的なPre-LayerNorm Transformerブロック"""
def __init__(self, d_model: int, n_heads: int, d_ff: int):
super().__init__()
self.ln1 = nn.LayerNorm(d_model)
self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
self.ln2 = nn.LayerNorm(d_model)
self.mlp = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Linear(d_ff, d_model),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
norm_x = self.ln1(x)
attn_out, _ = self.attn(norm_x, norm_x, norm_x)
x = x + attn_out
x = x + self.mlp(self.ln2(x))
return x
class LoopedReasoningTransformer(nn.Module):
"""
重み共有ループとKLダイバージェンス適応的早期終了を備えたLLM
"""
def __init__(
self,
vocab_size: int = 1000,
d_model: int = 256,
n_heads: int = 4,
d_ff: int = 1024,
n_shared_layers: int = 4,
max_loops: int = 4,
):
super().__init__()
self.d_model = d_model
self.max_loops = max_loops
self.embed = nn.Embedding(vocab_size, d_model)
# 重み共有されるコアスタック
self.shared_stack = nn.ModuleList([
TransformerBlock(d_model, n_heads, d_ff)
for _ in range(n_shared_layers)
])
# 入力状態と再帰状態を統合する射影層
self.latent_proj = nn.Linear(d_model * 2, d_model)
self.final_ln = nn.LayerNorm(d_model)
self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
def forward(
self,
input_ids: torch.Tensor,
kl_threshold: float = 0.02,
) -> tuple[torch.Tensor, int]:
"""
順伝播処理
Args:
input_ids: (batch_size, seq_len)
kl_threshold: 推論時早期終了のKLダイバージェンスしきい値
Returns:
logits: (batch_size, seq_len, vocab_size)
executed_loops: 実行されたループパス数
"""
batch_size, seq_len = input_ids.shape
h_0 = self.embed(input_ids)
h = h_0
prev_probs = None
executed_loops = 0
for loop_idx in range(self.max_loops):
executed_loops += 1
# パス2以降は初期入力表現を結合して射影
if loop_idx > 0:
h = self.latent_proj(torch.cat([h, h_0], dim=-1))
# 共有スタックの実行
for layer in self.shared_stack:
h = layer(h)
# 現在パスの語彙分布を算出
current_logits = self.lm_head(self.final_ln(h))
current_probs = F.softmax(current_logits, dim=-1)
# 推論時の適応的早期終了判定
if not self.training and prev_probs is not None:
log_probs = F.log_softmax(current_logits, dim=-1)
# バッチ全体の平均KLダイバージェンスを計算
kl_div = F.kl_div(
log_probs,
prev_probs,
reduction="batchmean",
log_target=False,
)
if kl_div.item() < kl_threshold:
break
prev_probs = current_probs.detach()
return current_logits, executed_loops
# 動作検証スクリプト
if __name__ == "__main__":
torch.manual_seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = LoopedReasoningTransformer(
vocab_size=1000,
d_model=128,
n_heads=4,
d_ff=512,
n_shared_layers=3,
max_loops=4,
).to(device)
dummy_input = torch.randint(0, 1000, (2, 16)).to(device)
# 訓練モード(全ループ実行)
model.train()
logits, loops = model(dummy_input)
print(f"[Training] Output shape: {logits.shape}, Loops executed: {loops}")
# 評価モード(KL収束判定による早期終了)
model.eval()
with torch.no_grad():
logits_eval, loops_eval = model(dummy_input, kl_threshold=0.08)
print(f"[Inference] Output shape: {logits_eval.shape}, Loops executed: {loops_eval}")
ベンチマークと実務でのトレードオフ
Looped Transformerは万能の解決策ではなく、計算量とメモリ配置において明確なトレードオフが存在する。
| 評価項目 | Standard (24L) | Deep (48L) | Looped (24L x 2) | SMELT (MoE Looped) |
|---|---|---|---|---|
| 重みパラメータ数 | 基準 (1.0x) | 2.0x | 1.0x (共有) | 2.0x (MoEで回復) |
| トークンあたり計算量 (FLOPs) | 基準 (1.0x) | 2.0x | 2.0x | 1.0x (幅縮小で適合) |
| KVキャッシュ容量 (VRAM) | 基準 (1.0x) | 2.0x | 2.0x (独立保持必須) | 1.0x (設計適合) |
| 訓練時活性化値メモリ | 基準 (1.0x) | 2.0x | 2.0x (アンロール) | 1.5x |
| 事実知識記憶容量 | 基準 (1.0x) | 2.0x | 1.0x (不変) | 1.8x |
| 多段階論理推論精度 | 基準 | 高い | 高い | 極めて高い |
| 生成レイテンシ (固定パス時) | 基準 (1.0x) | 2.0x | 2.0x | 1.0〜1.2x |
本番運用における4つの技術的課題(Gotchas)
- KVキャッシュ共有の誤謬: 重みを共有しているからといって、パス1とパス2で同一のKVキャッシュを再利用することは不可能である。同一トークンであってもパス1とパス2では入力される隠れ状態ベクトルが異なり、生成されるKey/Value表現も全く異なる。Nanbeigeの実験でもKVキャッシュ共有は致命的なパープレキシティ悪化を招くことが報告されており、実効深度に比例した独立KVキャッシュメモリの確保が必須となる。
- 逆伝播活性化メモリの急増: 学習時、勾配はアンロールされた全反復パスを逆伝播する。重みパラメータが半分であっても、保持すべき活性化テンソル(Activation Tensor)は44層分存在する。Activation Checkpointing(勾配再計算)を適切に構成しなければ、Out-of-Memory(OOM)を引き起こす原因となる。
- 量子化誤差の累積増幅: 同一の重みテンソルが1つのフォワードパス内で複数回適用されるため、FP8やINT4などの量子化ノイズが反復ごとに非線形に累積する。低ビット量子化を適用する際は、重み共有ブロックに対する外れ値(Outlier)抑制キャリブレーションを厳格に行う必要がある。
- バッチ推論時のStraggler遅延: トークン単位またはシーケンス単位で適応的早期終了を適用する場合、同一バッチ内に4ループを要する難問と1ループで完了する平易な問題が混在すると、GPUの同期処理によってバッチ全体の推論時間は最遅トークンに引きずられる。実サービング環境では、動的ルーティングに応じたバッチ再編成(Dynamic Micro-batching)インフラが不可欠となる。
導入チェックリストと今後の検証ステップ
- アーキテクチャ初期設計: 既存モデルのアップサイクリングではなく、スクラッチからの事前学習を前提にパイプラインを組む。パス数は2を基準とし、3パス以上は収穫逓減を厳格にプロファイリングする。
- KVキャッシュ管理機構の実装: vLLMやTensorRT-LLM等のサービングフレームワークを改変し、共有スタックの反復ごとに独立したレイヤーIDとしてKVキャッシュブロックを割り当てるルーティングロジックを構成する。
- Activation Checkpointingの最適化: PyTorchの
torch.utils.checkpointを再帰ループ単位で適用し、フォワードパスの活性化値を保存せず再計算することでVRAMフットプリントを抑制する。 - KLダイバージェンスしきい値のスイートスポット探索: タスクセット(MMLU等の知識系 vs GSM8K等の推論系)ごとに早期終了しきい値と精度のトレードオフ曲線をプロットし、プロダクション向けデフォルト値を確定する。
参考文献・参照リソース
- Sebastian Raschka: GPT-6 Astra, Looped Transformers, and Hidden Reasoning (Ahead of AI, Sep 2026)
- SMELT研究チーム: SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers (arXiv:2609.01343, Sep 2026)
- Nanbeigeチーム: Nanbeige 4.2: Technical Report on Recurrent Depth Language Models (arXiv:2607.22083, Jul 2026)
- Zhu et al.: Beyond Parameters: Exploring Virtual Logic Depth for Scaling Laws (arXiv:2506.18233, Jun 2025)
- Geiping et al.: Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach (arXiv:2502.05171, Feb 2025)
- Dehghani et al.: Universal Transformers (ICLR 2019 / arXiv:1807.03819, 2018)
- The Information: Secret Technique Behind OpenAI's Astra Model Sparks Security Concerns (Sep 2026)