final model (main)

config.json

@@ -0,0 +1,36 @@
+{
+  "architectures": [
+    "LoopedForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_babyloop.BabyloopConfig",
+    "AutoModelForCausalLM": "modeling_babyloop.LoopedForCausalLM"
+  },
+  "bias": false,
+  "bos_token_id": 1,
+  "d_model": 768,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "ffn_hidden": 2048,
+  "fusion": null,
+  "hidden_size": 768,
+  "inject_input": true,
+  "k": 4,
+  "max_position_embeddings": 512,
+  "max_seq_len": 512,
+  "model_type": "babyloop",
+  "n_coda": 0,
+  "n_core": 12,
+  "n_heads": 12,
+  "n_layers": 12,
+  "n_prelude": 0,
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "rms_eps": 1e-05,
+  "rope_base": 10000.0,
+  "tie_embeddings": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.51.3",
+  "vocab_size": 16000
+}

configuration_babyloop.py

@@ -0,0 +1,83 @@
+"""BabyloopConfig — HF ``PretrainedConfig`` 互換のモデル設定。
+
+このファイルは ``save_pretrained`` 時に checkpoint ディレクトリへ複製され、
+公式 evaluation-pipeline 側（別プロセス）から ``trust_remote_code=True`` で
+import される。そのため **transformers と標準ライブラリ以外に依存しない**
+（``babyloop`` パッケージ内部を import しない）こと。
+
+フィールドと configs/model/*.yaml の対応は docs/architecture.md §6 を参照。
+"""
+
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class BabyloopConfig(PretrainedConfig):
+    """ループドTransformer（K=1で標準に縮退）の設定。
+
+    軸A（標準 vs 再帰）は ``k`` のみで切り替える単一実装。視覚（②④）の口も
+    持つが、テキストのみ（①③）では ``fusion=None`` で無効。
+    """
+
+    model_type = "babyloop"
+
+    def __init__(
+        self,
+        d_model: int = 768,
+        n_layers: int = 12,
+        n_heads: int = 12,
+        ffn_hidden: int = 2048,
+        n_prelude: int = 0,
+        n_core: int = 12,
+        n_coda: int = 0,
+        k: int = 1,
+        inject_input: bool = False,
+        vocab_size: int = 16000,
+        max_seq_len: int = 512,
+        rope_base: float = 10000.0,
+        rms_eps: float = 1e-5,
+        tie_embeddings: bool = True,
+        bias: bool = False,
+        dropout: float = 0.0,
+        fusion: str | None = None,
+        pad_token_id: int | None = None,
+        bos_token_id: int | None = None,
+        eos_token_id: int | None = None,
+        **kwargs,
+    ):
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.ffn_hidden = ffn_hidden
+        self.n_prelude = n_prelude
+        self.n_core = n_core
+        self.n_coda = n_coda
+        self.k = k
+        self.inject_input = inject_input
+        self.vocab_size = vocab_size
+        self.max_seq_len = max_seq_len
+        self.rope_base = rope_base
+        self.rms_eps = rms_eps
+        self.tie_embeddings = tie_embeddings
+        self.bias = bias
+        self.dropout = dropout
+        self.fusion = fusion
+
+        # HF 汎用ユーティリティが参照するエイリアス。
+        self.hidden_size = d_model
+        self.num_attention_heads = n_heads
+        self.num_hidden_layers = n_layers
+        self.max_position_embeddings = max_seq_len
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_embeddings,
+            **kwargs,
+        )
+
+    @property
+    def head_dim(self) -> int:
+        return self.d_model // self.n_heads

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e01d999d493ddf35b6090074ec6b77a3e0fccb455726d2f7e84be83431145ae9
+size 388976400

modeling_babyloop.py

@@ -0,0 +1,291 @@
+"""ループドTransformer本体（HF ``PreTrainedModel`` 互換・単一実装）。
+
+設計 docs/architecture.md §2 の Llama 系レシピ（RMSNorm / RoPE / SwiGLU /
+bias なし / weight tying）を素の PyTorch で実装。``forward`` は K 回ループする
+形で書き、``k=1`` で標準Transformerに厳密に縮退する（``tests/test_k1_equivalence.py``）。
+
+このファイルは ``save_pretrained`` 時に checkpoint へ複製され、公式
+evaluation-pipeline（別プロセス・``trust_remote_code=True``）から import される。
+そのため **torch / transformers / 標準ライブラリ以外に依存しない**こと。
+probing 用のループ毎中間表現は HF 標準 ``hidden_states``（層ごと）と混ぜず、
+別フィールド ``loop_hidden_states`` に格納する。
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers.modeling_outputs import ModelOutput
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_babyloop import BabyloopConfig
+
+
+# --- ビルディングブロック（Llama系レシピ）---------------------------------
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x.to(dtype)) * self.weight
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    # x: (B, H, T, head_dim); cos/sin: (1, 1, T, head_dim)
+    return x * cos + _rotate_half(x) * sin
+
+
+class Attention(nn.Module):
+    """RoPE 付き causal Multi-Head Attention（dropout なし）。"""
+
+    def __init__(self, config: BabyloopConfig):
+        super().__init__()
+        self.n_heads = config.n_heads
+        self.head_dim = config.d_model // config.n_heads
+        self.qkv_proj = nn.Linear(config.d_model, 3 * config.d_model, bias=config.bias)
+        self.o_proj = nn.Linear(config.d_model, config.d_model, bias=config.bias)
+
+    def forward(self, x, cos, sin, attn_bias):
+        B, T, C = x.shape
+        q, k, v = self.qkv_proj(x).split(C, dim=-1)
+        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        q = _apply_rope(q, cos, sin)
+        k = _apply_rope(k, cos, sin)
+        out = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=attn_bias, is_causal=attn_bias is None
+        )
+        out = out.transpose(1, 2).reshape(B, T, C)
+        return self.o_proj(out)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: BabyloopConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.d_model, config.ffn_hidden, bias=config.bias)
+        self.up_proj = nn.Linear(config.d_model, config.ffn_hidden, bias=config.bias)
+        self.down_proj = nn.Linear(config.ffn_hidden, config.d_model, bias=config.bias)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Block(nn.Module):
+    """Pre-Norm 残差ブロック: h += Attn(RMSNorm(h)); h += SwiGLU(RMSNorm(h))。"""
+
+    def __init__(self, config: BabyloopConfig):
+        super().__init__()
+        self.attn_norm = RMSNorm(config.d_model, config.rms_eps)
+        self.attn = Attention(config)
+        self.mlp_norm = RMSNorm(config.d_model, config.rms_eps)
+        self.mlp = SwiGLU(config)
+
+    def forward(self, h, cos, sin, attn_bias):
+        h = h + self.attn(self.attn_norm(h), cos, sin, attn_bias)
+        h = h + self.mlp(self.mlp_norm(h))
+        return h
+
+
+# --- 出力コンテナ -----------------------------------------------------------
+
+
+@dataclass
+class LoopedModelOutput(ModelOutput):
+    last_hidden_state: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    loop_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+@dataclass
+class LoopedCausalLMOutput(ModelOutput):
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    loop_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+# --- PreTrainedModel ラッパ -------------------------------------------------
+
+
+class LoopedPreTrainedModel(PreTrainedModel):
+    config_class = BabyloopConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        std = 0.02
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, RMSNorm):
+            nn.init.ones_(module.weight)
+
+
+class LoopedModel(LoopedPreTrainedModel):
+    """重み共有 core ブロックを K 回反復するバックボーン（lm_head なし）。"""
+
+    def __init__(self, config: BabyloopConfig):
+        super().__init__(config)
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
+        self.prelude = nn.ModuleList(Block(config) for _ in range(config.n_prelude))
+        self.core = nn.ModuleList(Block(config) for _ in range(config.n_core))
+        self.coda = nn.ModuleList(Block(config) for _ in range(config.n_coda))
+        self.final_norm = RMSNorm(config.d_model, config.rms_eps)
+
+        head_dim = config.d_model // config.n_heads
+        inv_freq = 1.0 / (
+            config.rope_base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
+        )
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def _rope(self, T: int, device, dtype):
+        t = torch.arange(T, device=device, dtype=torch.float32)
+        freqs = torch.outer(t, self.inv_freq.to(device))
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb.cos().to(dtype)[None, None], emb.sin().to(dtype)[None, None]
+
+    def _attn_bias(self, attention_mask, T, device, dtype):
+        # padding が無ければ None を返し、SDPA の is_causal 経路（高速）に乗せる。
+        if attention_mask is None or bool((attention_mask == 1).all()):
+            return None
+        causal = torch.ones(T, T, device=device, dtype=torch.bool).triu(1)
+        key_pad = attention_mask.to(device) == 0  # (B, T)
+        mask = causal[None, None] | key_pad[:, None, None, :]
+        bias = torch.zeros(mask.shape, device=device, dtype=dtype)
+        return bias.masked_fill(mask, torch.finfo(dtype).min)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        output_hidden_states=False,
+        **kwargs,
+    ) -> LoopedModelOutput:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        h = inputs_embeds
+        residual_input = inputs_embeds  # inject_input 用（③で本格化、①は false）
+        B, T, _ = h.shape
+        cos, sin = self._rope(T, h.device, h.dtype)
+        attn_bias = self._attn_bias(attention_mask, T, h.device, h.dtype)
+
+        all_hidden = [h] if output_hidden_states else None
+        loop_hidden = []
+
+        for block in self.prelude:
+            h = block(h, cos, sin, attn_bias)
+            if output_hidden_states:
+                all_hidden.append(h)
+        for _ in range(self.config.k):
+            for block in self.core:
+                h = block(h, cos, sin, attn_bias)
+                if output_hidden_states:
+                    all_hidden.append(h)
+            if self.config.inject_input:
+                h = h + residual_input
+            loop_hidden.append(h)
+        for block in self.coda:
+            h = block(h, cos, sin, attn_bias)
+            if output_hidden_states:
+                all_hidden.append(h)
+
+        h = self.final_norm(h)
+        return LoopedModelOutput(
+            last_hidden_state=h,
+            hidden_states=tuple(all_hidden) if output_hidden_states else None,
+            loop_hidden_states=tuple(loop_hidden),
+        )
+
+
+class LoopedForCausalLM(LoopedPreTrainedModel):
+    """言語モデリングヘッド付き（``AutoModelForCausalLM`` 互換）。"""
+
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: BabyloopConfig):
+        super().__init__(config)
+        self.model = LoopedModel(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new):
+        self.lm_head = new
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_hidden_states=False,
+        **kwargs,
+    ) -> LoopedCausalLMOutput:
+        out = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+        )
+        logits = self.lm_head(out.last_hidden_state)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        return LoopedCausalLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=out.hidden_states,
+            loop_hidden_states=out.loop_hidden_states,
+        )
+
+
+# 公式 evaluation-pipeline が trust_remote_code で AutoModel 系から読めるよう登録。
+# save_pretrained 時に auto_map と本ファイル群が checkpoint へ複製される。
+BabyloopConfig.register_for_auto_class()
+LoopedModel.register_for_auto_class("AutoModel")
+LoopedForCausalLM.register_for_auto_class("AutoModelForCausalLM")

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|bos|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|eos|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|pad|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|unk|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,44 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|bos|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|eos|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<|unk|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|bos|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|eos|>",
+  "extra_special_tokens": {},
+  "model_max_length": 512,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "PreTrainedTokenizer",
+  "unk_token": "<|unk|>"
+}