[UPDATE] Replace weights with the V5 model; update paper, model card and evaluation results

README.md

@@ -23,15 +23,15 @@ metrics:
 
 ## 1. Model Overview
 
-**Character-Based Transformer (CBT)** is an 18.95M-parameter character-level Transformer whose tokens carry a rule-assigned two-dimensional (2D) coordinate $(p, \tau)$ used by **Decoupled 2D Rotary Position Embeddings (2D-RoPE)** alongside **Weight-Tied Recurrent Depth ($R = 3$)**.
+**Character-Based Transformer (CBT-18.95M)** is an 18.94M-parameter character-level Transformer (111-token vocabulary) whose tokens carry a rule-assigned two-dimensional coordinate $(p, \tau)$ used by **Decoupled 2D Rotary Position Embeddings (2D-RoPE)**, with weight-tied recurrent depth ($R = 3$). This release is the **V5 model** (`champion_final.pt`): pretraining, search-copy and all-tools fine-tuning, V4 fine-tuning, then a 23,927-step V5 fine-tuning on a decontaminated online stream (35% search copy, 20% arithmetic replay of 1-25 digits, the rest multi-tool and chat). Coordinate shift masking uses $k=45$, so $p_{\max} = 25 + 45 = 70$.
 
-Trained on 1- to 25-digit arithmetic replay, multi-tool dialogues, and search copy, with random coordinate shift masking $k=45$ (macro column coordinates up to $p_{\max} = 25 + 45 = 70$), **CBT-18.95M** (`champion_final.pt`) achieves:
-- **100.0% (240/240; 95% Wilson CI [98.4, 100.0]%)** on random multi-carry sums from 16 to 50 digits, and **100.0% (18/18)** on structured 16- to 50-digit sums (including 50-digit cascading carry $99\dots99 + 1$). Digits 16–25 are seen in training; 30–50 digits are unseen.
-- **98.9% (89/90)** on in-distribution sums at 8, 12, and 15 digits (the single miss is at 15 digits).
-- **100.0% (20/20)** at 55 digits, then **70.0% (14/20)** at 60 digits, **40.0% (6/15)** at 70 digits, and **0.0%** at 80–100 digits (random sweep). A separate teacher-forced suite gives 80% at both 60 and 70 digits.
-- **Equalized 1–15d, $T=1024$ Control Study (12,800 Random Problems)**: When four equalized 100k models (`1D-RoPE Standard`, `1D-RoPE + Offset`, `1D-RoPE + Ruoss et al. 2023 Randomized Positions`, and `Decoupled 2D-RoPE k=32`) are trained on the exact same 1–15d dataset with $T=1024$ and evaluated on 200 random problems per length up to 100d, all three 1D/Ruoss variants score **0.0% (0/200, CI [0.0, 1.9]%)** across 16–100d, whereas equalized 2D-RoPE achieves **81.4% (977/1,200, CI [79.1, 83.5]%)** across 16–40d (up to 92.0% at 16d and 25d).
-- **100.0% (90/90)** on a 6-category multi-tool benchmark (the check counts tags only, not argument correctness); **80.8% (21/26)** tool routing and **50.0% (13/26)** exact argument extraction on a 26-case phrasing suite with audited entities.
-- **100.0% (18/18)** zero-shot OOD tool routing, with **55.6% (10/18)** exact argument copying on 18 fresh prompts.
+Results (greedy decoding; Wilson 95% CIs are in the paper and in `results/`):
+- **Addition**, 30 random problems per length (70% uniform, 15% dense carry, 15% cascades), scored as teacher-forced greedy agreement with the gold scratchpad and answer: **150/150** at 4-16 digits, **60/60** at 20 and 25 digits (seen), **210/210** at 26-55 digits (unseen), **25/30** at 60 and **19/30** at 70 digits. Pooled over the unseen 26-70 digits: **254/270**.
+- **Tool use**: **26/26** exact calls on a 26-prompt phrasing suite, **18/18** on 18 out-of-distribution prompts, and **30/30** exact copies of entities absent from the training corpus (unseen-30). The V5 data was broadened after the previous model's failures on the first two suites were seen, so the unseen-30 set is the cleaner evidence.
+- **1D-RoPE control**: a twin with the same architecture, V5 data and step count but standard 1D-RoPE scores **0/270** at 26-70 digits and **73/74** exact tool calls. The 2D coordinates account for the arithmetic extrapolation, not for the tool-use results.
+- **Baselines**: standard 1D-RoPE models trained on 1-12 digit sums score 0/30 at every length from 14 to 35 digits (character and BPE vocabularies; looped, hierarchical and hybrid variants too).
+
+All results are single training runs. Details, limitations and the failure analysis are in the paper.
 
 ---
 
@@ -39,61 +39,37 @@ Trained on 1- to 25-digit arithmetic replay, multi-tool dialogues, and search co
 
 | Parameter | Specification |
 | :--- | :--- |
-| **Model Name** | Character-Based Transformer (CBT-18.95M) |
-| **Total Parameters** | 18,946,560 (~18.95M) |
-| **Non-Embedding Parameters** | 18,887,680 (~18.89M) |
+| **Total Parameters** | 18,944,512 |
+| **Non-Embedding Parameters** | 18,887,680 |
 | **Physical Layers (L)** | 6 |
-| **Recurrent Passes (R)** | 3 (Effective Depth = 18 layers) |
-| **Hidden Dimension (d_model)** | 512 |
-| **Attention Heads (H)** | 8 (Head dimension = 64) |
-| **Max Sequence Length (T)** | 1024 (Coordinate cache up to 4096) |
-| **Positional Embedding** | Decoupled 2D-RoPE (d/2 macro column $p$, d/2 micro step $\tau$; rule-assigned via `compute_2d_coords`) |
-| **Coordinate Shift Masking (k)** | $k = 45$ ($p_{\max} = 25 + 45 = 70$ during 1–25d arithmetic replay) |
-| **Vocabulary Size** | 115 (14 structural control tags + 101 ASCII characters) |
-| **Weight Tying** | W_te == W_lm_head |
-
----
-
-## 3. Empirical Evaluation Benchmarks (Greedy Decoding)
-
-| Benchmark Suite | Scope / Test Cases | Final Score | 95% Wilson CI |
-| :--- | :--- | :--- | :--- |
-| **In-Distribution Sums (8, 12, 15 Digits)** | Seen in Training | **98.9% (89/90)** | [94.0, 99.8]% |
-| **16–50 Digit Structured Sums** | 18 Cases (incl. $99\dots99 + 1$) | **100.0% (18/18)** | [82.4, 100.0]% |
-| **16–50 Digit Random Multi-Carry Sweep** | 8 Lengths $\times$ 30 Cases ($n=240$) | **100.0% (240/240)** | [98.4, 100.0]% |
-| **55–70 Digit Random Sweep** | 55d: 20 cases; 60d: 20; 70d: 15 | 55d: **100.0%**; 60d: **70.0%**; 70d: **40.0%** | 55d [83.9, 100.0]%; 60d [48.1, 85.5]%; 70d [19.8, 64.3]% |
-| **80–100 Digit Random Sweep** | 10 Cases per length | **0.0%** | [0.0, 27.8]% |
-| **Multi-Tool Benchmark (6 Categories)** | 90 Dialogues (tag-presence check) | **100.0% (90/90)** | [95.9, 100.0]% |
-| **Multi-Tool Phrasing Routing** | 26 Audited Cases, 9 Schemas | **80.8% (21/26)** | [62.1, 91.5]% |
-| **Multi-Tool Phrasing Exact Arguments** | 26 Audited Cases | **50.0% (13/26)** | [32.1, 67.9]% |
-| **Zero-Shot OOD Routing** | 18 Fresh Prompts | **100.0% (18/18)** | [82.4, 100.0]% |
-| **Zero-Shot OOD Exact Argument Copy** | 18 Fresh Prompts | **55.6% (10/18)** | [33.7, 75.4]% |
+| **Recurrent Passes (R)** | 3 (effective depth 18) |
+| **Hidden Dimension** | 512 |
+| **Attention Heads** | 8 (head dimension 64) |
+| **Context** | 1024 tokens |
+| **Positional Embedding** | Decoupled 2D-RoPE (half of each head rotated by place $p$, half by step $\tau$; rule-assigned) |
+| **Coordinate Shift Masking** | $k = 45$ |
+| **Vocabulary** | 111 (14 control tokens + 97 characters) |
+| **Weight Tying** | input embedding = output projection |
 
 ---
 
-## 4. Usage & Inference
+## 3. Usage
 
-### Local Quickstart
+The file is a `torch.save` dictionary with the config and `model_state_dict`.
 
 ```python
 import torch
-from model import ScaledGPT, ScaledConfig
+from model import ScaledGPT
 from tokenizer import ScaledTokenizer
 
-# 1. Load configuration and model
-config = ScaledConfig()
-model = ScaledGPT(config)
-
-# 2. Load weights
-state_dict = torch.load("champion_final.pt", map_location="cpu", weights_only=True)
-model.load_state_dict(state_dict)
-model.eval()
-
-# 3. Initialize tokenizer
+ckpt = torch.load("champion_final.pt", map_location="cpu", weights_only=False)
 tokenizer = ScaledTokenizer()
+model = ScaledGPT(ckpt["config"], tokenizer=tokenizer)
+model.load_state_dict(ckpt["model_state_dict"])
+model.eval()
 ```
 
-### Interactive CLI
+Interactive CLI:
 
 ```bash
 python interactive.py --ckpt champion_final.pt
@@ -101,6 +77,12 @@ python interactive.py --ckpt champion_final.pt
 
 ---
 
+## 4. Files
+
+`champion_final.pt` (V5 model), source files for stand-alone execution, `RESEARCH_PAPER.md`, and the evaluation results in JSON.
+
+---
+
 ## 5. Citation
 
 ```bibtex

cbt1d_v5_arithmetic.json

@@ -0,0 +1,224 @@
+{
+  "checkpoint": "champion_model/checkpoints/1d/cbt_18m_1d_rope_R3_k45_sftv5.zip",
+  "max_train_operand_digits": 25,
+  "arithmetic": {
+    "per_length": {
+      "20": {
+        "digits": 20,
+        "mode": "teacher_forced",
+        "seen_in_training": true,
+        "n": 30,
+        "exact": 28,
+        "exact_ci": [
+          78.7,
+          98.2
+        ],
+        "think_exact": 28,
+        "think_ci": [
+          78.7,
+          98.2
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 9.7
+      },
+      "25": {
+        "digits": 25,
+        "mode": "teacher_forced",
+        "seen_in_training": true,
+        "n": 30,
+        "exact": 29,
+        "exact_ci": [
+          83.3,
+          99.4
+        ],
+        "think_exact": 29,
+        "think_ci": [
+          83.3,
+          99.4
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 12.08
+      },
+      "26": {
+        "digits": 26,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 0,
+        "exact_ci": [
+          0.0,
+          11.4
+        ],
+        "think_exact": 0,
+        "think_ci": [
+          0.0,
+          11.4
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 12.42
+      },
+      "30": {
+        "digits": 30,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 0,
+        "exact_ci": [
+          0.0,
+          11.4
+        ],
+        "think_exact": 0,
+        "think_ci": [
+          0.0,
+          11.4
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 14.01
+      },
+      "35": {
+        "digits": 35,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 0,
+        "exact_ci": [
+          0.0,
+          11.4
+        ],
+        "think_exact": 0,
+        "think_ci": [
+          0.0,
+          11.4
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 17.04
+      },
+      "40": {
+        "digits": 40,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 0,
+        "exact_ci": [
+          0.0,
+          11.4
+        ],
+        "think_exact": 0,
+        "think_ci": [
+          0.0,
+          11.4
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 19.7
+      },
+      "45": {
+        "digits": 45,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 0,
+        "exact_ci": [
+          0.0,
+          11.4
+        ],
+        "think_exact": 0,
+        "think_ci": [
+          0.0,
+          11.4
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 22.33
+      },
+      "50": {
+        "digits": 50,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 0,
+        "exact_ci": [
+          0.0,
+          11.4
+        ],
+        "think_exact": 0,
+        "think_ci": [
+          0.0,
+          11.4
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 23.61
+      },
+      "55": {
+        "digits": 55,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 0,
+        "exact_ci": [
+          0.0,
+          11.4
+        ],
+        "think_exact": 0,
+        "think_ci": [
+          0.0,
+          11.4
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 26.12
+      },
+      "60": {
+        "digits": 60,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 0,
+        "exact_ci": [
+          0.0,
+          11.4
+        ],
+        "think_exact": 0,
+        "think_ci": [
+          0.0,
+          11.4
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 28.45
+      },
+      "70": {
+        "digits": 70,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 0,
+        "exact_ci": [
+          0.0,
+          11.4
+        ],
+        "think_exact": 0,
+        "think_ci": [
+          0.0,
+          11.4
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 33.21
+      }
+    },
+    "summary": {
+      "unseen": {
+        "n": 270,
+        "exact": 0,
+        "exact_ci": [
+          0.0,
+          1.4
+        ]
+      },
+      "seen": {
+        "n": 60,
+        "exact": 57,
+        "exact_ci": [
+          86.3,
+          98.3
+        ]
+      }
+    }
+  }
+}

champion_final.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:634bf18ed4173bc73e58970c6e843fbb4ccec92fde2b12c01bbe2fe04fe85a13
-size 75800439
+oid sha256:c6fab36eaadbc402441691ded7ed9e26990f3718e6cefcb5a889691b7fc3b9e5
+size 75799812

champion_v3_arithmetic.json

@@ -0,0 +1,224 @@
+{
+  "checkpoint": "champion_model/checkpoints/champion_v3_final.zip",
+  "max_train_operand_digits": 25,
+  "arithmetic": {
+    "per_length": {
+      "20": {
+        "digits": 20,
+        "mode": "teacher_forced",
+        "seen_in_training": true,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 60.31
+      },
+      "25": {
+        "digits": 25,
+        "mode": "teacher_forced",
+        "seen_in_training": true,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 59.46
+      },
+      "26": {
+        "digits": 26,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 67.86
+      },
+      "30": {
+        "digits": 30,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 63.91
+      },
+      "35": {
+        "digits": 35,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 76.18
+      },
+      "40": {
+        "digits": 40,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 84.75
+      },
+      "45": {
+        "digits": 45,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 91.2
+      },
+      "50": {
+        "digits": 50,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 115.78
+      },
+      "55": {
+        "digits": 55,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 107.83
+      },
+      "60": {
+        "digits": 60,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 25,
+        "exact_ci": [
+          66.4,
+          92.7
+        ],
+        "think_exact": 25,
+        "think_ci": [
+          66.4,
+          92.7
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 120.42
+      },
+      "70": {
+        "digits": 70,
+        "mode": "teacher_forced",
+        "seen_in_training": false,
+        "n": 30,
+        "exact": 19,
+        "exact_ci": [
+          45.5,
+          78.1
+        ],
+        "think_exact": 19,
+        "think_ci": [
+          45.5,
+          78.1
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 144.89
+      }
+    },
+    "summary": {
+      "unseen": {
+        "n": 270,
+        "exact": 254,
+        "exact_ci": [
+          90.6,
+          96.3
+        ]
+      },
+      "seen": {
+        "n": 60,
+        "exact": 60,
+        "exact_ci": [
+          94.0,
+          100.0
+        ]
+      }
+    }
+  }
+}

champion_v3_arithmetic_4_16d.json

@@ -0,0 +1,108 @@
+{
+  "checkpoint": "checkpoints/cbt/cbt_18m_2d_rope_R3_k45_sftv5.pt",
+  "max_train_operand_digits": 25,
+  "arithmetic": {
+    "per_length": {
+      "4": {
+        "digits": 4,
+        "mode": "teacher_forced",
+        "seen_in_training": true,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 3.29
+      },
+      "8": {
+        "digits": 8,
+        "mode": "teacher_forced",
+        "seen_in_training": true,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 5.23
+      },
+      "12": {
+        "digits": 12,
+        "mode": "teacher_forced",
+        "seen_in_training": true,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 6.68
+      },
+      "14": {
+        "digits": 14,
+        "mode": "teacher_forced",
+        "seen_in_training": true,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 7.93
+      },
+      "16": {
+        "digits": 16,
+        "mode": "teacher_forced",
+        "seen_in_training": true,
+        "n": 30,
+        "exact": 30,
+        "exact_ci": [
+          88.6,
+          100.0
+        ],
+        "think_exact": 30,
+        "think_ci": [
+          88.6,
+          100.0
+        ],
+        "gen_tokens": 0,
+        "elapsed_sec": 8.82
+      }
+    },
+    "summary": {
+      "seen": {
+        "n": 150,
+        "exact": 150,
+        "exact_ci": [
+          97.5,
+          100.0
+        ]
+      }
+    }
+  }
+}

champion_v3_ood.json

@@ -0,0 +1,17 @@
+{
+  "checkpoint": "champion_model/checkpoints/champion_v3_final.zip",
+  "max_train_operand_digits": 25,
+  "ood": {
+    "total": 18,
+    "routed": 18,
+    "routing_ci": [
+      82.4,
+      100.0
+    ],
+    "exact_matches": 18,
+    "exact_ci": [
+      82.4,
+      100.0
+    ]
+  }
+}

champion_v3_phrasing.json

@@ -0,0 +1,17 @@
+{
+  "checkpoint": "champion_model/checkpoints/champion_v3_final.zip",
+  "max_train_operand_digits": 25,
+  "phrasing": {
+    "total": 26,
+    "routing_passed": 26,
+    "routing_ci": [
+      87.1,
+      100.0
+    ],
+    "exact_args_passed": 26,
+    "exact_args_ci": [
+      87.1,
+      100.0
+    ]
+  }
+}

champion_v3_unseen_ood.json

@@ -0,0 +1,354 @@
+{
+  "total": 30,
+  "routing_passed": 30,
+  "routing_pct": 100.0,
+  "routing_ci": [
+    88.6,
+    100.0
+  ],
+  "exact_copy_passed": 30,
+  "headline_exact_pct": 100.0,
+  "headline_exact_ci": [
+    88.6,
+    100.0
+  ],
+  "copy_given_routed_pct": 100.0,
+  "copy_given_routed_ci": [
+    88.6,
+    100.0
+  ],
+  "binomial_p_value_vs_50pct": 9.313225746154785e-10,
+  "is_statistically_gt_50pct": true,
+  "itemized": [
+    {
+      "id": 1,
+      "domain": "Unseen Figure",
+      "prompt": "<USER>Who was Baruch Spinoza?</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "Baruch Spinoza",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"Baruch Spinoza\")</CALL>",
+      "parsed_call": "search(query=\"Baruch Spinoza\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 2,
+      "domain": "Unseen Figure",
+      "prompt": "<USER>Who was Hypatia of Alexandria?</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "Hypatia of Alexandria",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"Hypatia of Alexandria\")</CALL>",
+      "parsed_call": "search(query=\"Hypatia of Alexandria\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 3,
+      "domain": "Unseen Figure",
+      "prompt": "<USER>Who was Ludwig van Beethoven?</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "Ludwig van Beethoven",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"Ludwig van Beethoven\")</CALL>",
+      "parsed_call": "search(query=\"Ludwig van Beethoven\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 4,
+      "domain": "Unseen Figure",
+      "prompt": "<USER>Who was Johannes Vermeer?</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "Johannes Vermeer",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"Johannes Vermeer\")</CALL>",
+      "parsed_call": "search(query=\"Johannes Vermeer\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 5,
+      "domain": "Unseen Figure",
+      "prompt": "<USER>Who is Ada Yonath?</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "Ada Yonath",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"Ada Yonath\")</CALL>",
+      "parsed_call": "search(query=\"Ada Yonath\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 6,
+      "domain": "Unseen Figure",
+      "prompt": "<USER>Who was Srinivasa Ramanujan?</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "Srinivasa Ramanujan",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"Srinivasa Ramanujan\")</CALL>",
+      "parsed_call": "search(query=\"Srinivasa Ramanujan\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 7,
+      "domain": "Unseen City",
+      "prompt": "<USER>What is the weather in Addis Ababa?</USER><THINK>",
+      "expected_tool": "get_weather",
+      "expected_arg": "Addis Ababa",
+      "raw_response": "The user requests current weather information. I will query the get_weather tool.</THINK><CALL>get_weather(city=\"Addis Ababa\")</CALL>",
+      "parsed_call": "get_weather(city=\"Addis Ababa\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 8,
+      "domain": "Unseen City",
+      "prompt": "<USER>What is the weather in Ulaanbaatar?</USER><THINK>",
+      "expected_tool": "get_weather",
+      "expected_arg": "Ulaanbaatar",
+      "raw_response": "The user requests current weather information. I will query the get_weather tool.</THINK><CALL>get_weather(city=\"Ulaanbaatar\")</CALL>",
+      "parsed_call": "get_weather(city=\"Ulaanbaatar\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 9,
+      "domain": "Unseen City",
+      "prompt": "<USER>Weather on Antananarivo</USER><THINK>",
+      "expected_tool": "get_weather",
+      "expected_arg": "Antananarivo",
+      "raw_response": "The user is asking for weather conditions. Querying get_weather.</THINK><CALL>get_weather(city=\"Antananarivo\")</CALL>",
+      "parsed_call": "get_weather(city=\"Antananarivo\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 10,
+      "domain": "Unseen City",
+      "prompt": "<USER>How is the weather in Dar es Salaam</USER><THINK>",
+      "expected_tool": "get_weather",
+      "expected_arg": "Dar es Salaam",
+      "raw_response": "The user is asking for weather conditions. Querying get_weather.</THINK><CALL>get_weather(city=\"Dar es Salaam\")</CALL>",
+      "parsed_call": "get_weather(city=\"Dar es Salaam\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 11,
+      "domain": "Unseen City",
+      "prompt": "<USER>Check the weather for Belo Horizonte</USER><THINK>",
+      "expected_tool": "get_weather",
+      "expected_arg": "Belo Horizonte",
+      "raw_response": "The user is asking for weather conditions. Querying get_weather.</THINK><CALL>get_weather(city=\"Belo Horizonte\")</CALL>",
+      "parsed_call": "get_weather(city=\"Belo Horizonte\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 12,
+      "domain": "Unseen City",
+      "prompt": "<USER>Forecast for Chiang Mai.</USER><THINK>",
+      "expected_tool": "get_weather",
+      "expected_arg": "Chiang Mai",
+      "raw_response": "The user is asking for weather conditions. Querying get_weather.</THINK><CALL>get_weather(city=\"Chiang Mai\")</CALL>",
+      "parsed_call": "get_weather(city=\"Chiang Mai\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 13,
+      "domain": "Unseen Science",
+      "prompt": "<USER>What are quarks and leptons</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "quarks and leptons",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"quarks and leptons\")</CALL>",
+      "parsed_call": "search(query=\"quarks and leptons\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 14,
+      "domain": "Unseen Science",
+      "prompt": "<USER>Define epigenetics</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "epigenetics",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"epigenetics\")</CALL>",
+      "parsed_call": "search(query=\"epigenetics\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 15,
+      "domain": "Unseen Science",
+      "prompt": "<USER>Explain superconductivity</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "superconductivity",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"superconductivity\")</CALL>",
+      "parsed_call": "search(query=\"superconductivity\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 16,
+      "domain": "Unseen Science",
+      "prompt": "<USER>What is the mitochondrial matrix</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "mitochondrial matrix",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"mitochondrial matrix\")</CALL>",
+      "parsed_call": "search(query=\"mitochondrial matrix\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 17,
+      "domain": "Unseen Science",
+      "prompt": "<USER>Explain bose einstein condensate</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "bose einstein condensate",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"bose einstein condensate\")</CALL>",
+      "parsed_call": "search(query=\"bose einstein condensate\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 18,
+      "domain": "Unseen Science",
+      "prompt": "<USER>What is asymptotic freedom?</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "asymptotic freedom",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"asymptotic freedom\")</CALL>",
+      "parsed_call": "search(query=\"asymptotic freedom\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 19,
+      "domain": "Unseen Concept",
+      "prompt": "<USER>What is a croissant?</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "croissant",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"croissant\")</CALL>",
+      "parsed_call": "search(query=\"croissant\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 20,
+      "domain": "Unseen Concept",
+      "prompt": "<USER>Define baklava</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "baklava",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"baklava\")</CALL>",
+      "parsed_call": "search(query=\"baklava\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 21,
+      "domain": "Unseen Concept",
+      "prompt": "<USER>What is shakshuka?</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "shakshuka",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"shakshuka\")</CALL>",
+      "parsed_call": "search(query=\"shakshuka\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 22,
+      "domain": "Unseen Concept",
+      "prompt": "<USER>Tell me about the astrolabe</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "astrolabe",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"astrolabe\")</CALL>",
+      "parsed_call": "search(query=\"astrolabe\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 23,
+      "domain": "Unseen Concept",
+      "prompt": "<USER>What is a clavichord</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "clavichord",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"clavichord\")</CALL>",
+      "parsed_call": "search(query=\"clavichord\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 24,
+      "domain": "Unseen Concept",
+      "prompt": "<USER>Explain origami</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "origami",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"origami\")</CALL>",
+      "parsed_call": "search(query=\"origami\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 25,
+      "domain": "Unseen Nonce",
+      "prompt": "<USER>Who is Vornik Plaxen?</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "Vornik Plaxen",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"Vornik Plaxen\")</CALL>",
+      "parsed_call": "search(query=\"Vornik Plaxen\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 26,
+      "domain": "Unseen Nonce",
+      "prompt": "<USER>Define Zyntral Krelvok</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "Zyntral Krelvok",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"Zyntral Krelvok\")</CALL>",
+      "parsed_call": "search(query=\"Zyntral Krelvok\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 27,
+      "domain": "Unseen Nonce",
+      "prompt": "<USER>Weather on Brixendale</USER><THINK>",
+      "expected_tool": "get_weather",
+      "expected_arg": "Brixendale",
+      "raw_response": "The user is asking for weather conditions. Querying get_weather.</THINK><CALL>get_weather(city=\"Brixendale\")</CALL>",
+      "parsed_call": "get_weather(city=\"Brixendale\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 28,
+      "domain": "Unseen Nonce",
+      "prompt": "<USER>What is Glarndoxification?</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "Glarndoxification",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"Glarndoxification\")</CALL>",
+      "parsed_call": "search(query=\"Glarndoxification\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 29,
+      "domain": "Unseen Nonce",
+      "prompt": "<USER>What is a Trundleblump</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "Trundleblump",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"Trundleblump\")</CALL>",
+      "parsed_call": "search(query=\"Trundleblump\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    },
+    {
+      "id": 30,
+      "domain": "Unseen Nonce",
+      "prompt": "<USER>Tell me about a Skrevling</USER><THINK>",
+      "expected_tool": "search",
+      "expected_arg": "Skrevling",
+      "raw_response": "The user is asking a factual question. I will query the search tool.</THINK><CALL>search(query=\"Skrevling\")</CALL>",
+      "parsed_call": "search(query=\"Skrevling\")",
+      "route_ok": true,
+      "exact_copy_ok": true
+    }
+  ]
+}

interactive.py

@@ -242,7 +242,7 @@ def run_dialogue_prompt(model: ScaledGPT, tokenizer: ScaledTokenizer, user_input
 
 def main():
     parser = argparse.ArgumentParser(description="ScaledGPT Interactive CLI")
-    parser.add_argument("--ckpt", type=str, default="champion_model/scaled_t4/checkpoints/champion_final.pt", help="Path to checkpoint")
+    parser.add_argument("--ckpt", type=str, default="checkpoints/cbt/cbt_18m_2d_rope_R3_k45_sftv5.pt", help="Path to checkpoint")
     parser.add_argument("--prompt", type=str, default=None, help="Execute a single prompt and exit")
     parser.add_argument("--threads", type=int, default=8, help="Number of CPU threads for inference")
     args = parser.parse_args()
@@ -258,15 +258,10 @@ def main():
     candidates = [
         args.ckpt,
         os.path.join(ROOT_DIR, args.ckpt),
-        os.path.join(ROOT_DIR, "champion_model/scaled_t4/checkpoints/champion_final.pt"),
-        os.path.join(ROOT_DIR, "champion_model/checkpoints/champion_final.pt"),
-        os.path.join(ROOT_DIR, "champion_model/scaled_t4/checkpoints/scaled_t4_sft_all_tools.pt"),
-        os.path.join(ROOT_DIR, "champion_model/scaled_t4/checkpoints/scaled_t4_sft_search.pt"),
-        os.path.join(ROOT_DIR, "champion_model/scaled_t4/checkpoints/scaled_t4_final.pt"),
-        os.path.join(ROOT_DIR, "champion_model/scaled_t4/checkpoints/scaled_t4_final_timeout_step_45059.pt"),
-        os.path.join(ROOT_DIR, "champion_model/checkpoints/scaled_t4_sft_search_step_7000.zip"),
-        os.path.join(ROOT_DIR, "champion_model/checkpoints/scaled_t4_final.zip"),
-        os.path.join(ROOT_DIR, "champion_model/checkpoints/scaled_t4_final_timeout_step_45059.zip"),
+        os.path.join(ROOT_DIR, "checkpoints/cbt/cbt_18m_2d_rope_R3_k45_sftv5.pt"),
+        os.path.join(ROOT_DIR, "champion_final.pt"),
+        os.path.join(ROOT_DIR, "checkpoints/cbt/cbt_18m_2d_rope_R3_k32_base.pt"),
+        os.path.join(ROOT_DIR, "checkpoints/cbt/cbt_18m_2d_rope_R3_k32_pretrain.pt"),
     ]
     resolved_path = None
     for p in candidates:
@@ -278,8 +273,11 @@ def main():
         try:
             from huggingface_hub import hf_hub_download
             print("[INFO] Local weights not found. Automatically pulling champion weights from Hugging Face (Neithanvarvas/CBT)...")
-            target_dir = os.path.join(ROOT_DIR, "champion_model", "scaled_t4", "checkpoints")
+            target_dir = os.path.join(ROOT_DIR, "checkpoints", "cbt")
             resolved_path = hf_hub_download(repo_id="Neithanvarvas/CBT", filename="champion_final.pt", local_dir=target_dir)
+            local_path = os.path.join(ROOT_DIR, "checkpoints/cbt/cbt_18m_2d_rope_R3_k45_sftv5.pt")
+            os.replace(resolved_path, local_path)
+            resolved_path = local_path
             print(f"[PASS] Successfully pulled weights to: {resolved_path}")
         except Exception as e:
             print(f"[WARN] Automatic Hugging Face pull unsuccessful ({e}).")

model.py

@@ -8,8 +8,9 @@ Architecture:
 - Recurrent passes R = 3 (Effective depth = 18 layers at zero parameter overhead)
 - Residual recurrence update: x = x + (1/sqrt(R)) * (x_res - x)
 - Tied embedding and output projection weights (wte.weight == lm_head.weight)
-- Decoupled 2D-RoPE (Macro place p, Micro step tau)
-- Selective shift masking: (place_coords > 0).long()
+- Decoupled 2D-RoPE (Macro place p, Micro step tau); pos_emb_type="rope_1d" gives the
+  standard 1D-RoPE ablation (positions 0..T-1 over the full head_dim, same parameters)
+- Selective shift masking: (place_coords > 0).long() (2D only)
 - O(1) SDPA causal attention
 - Non-embedding params: ~18,887,680 | Total params: ~18,953,216 (~18.95M target)
 """
@@ -106,6 +107,19 @@ def apply_2d_rope(
     return torch.cat([rot_place, rot_step], dim=-1)
 
 
+def apply_1d_rope(x: torch.Tensor, pos: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Applies standard 1D-RoPE over the full head_dim of [B, n_head, T, head_dim] at positions pos [B, T]."""
+    pos = torch.clamp(pos, 0, cos.size(0) - 1)
+    c = cos[pos].unsqueeze(1).to(dtype=x.dtype, device=x.device)
+    s = sin[pos].unsqueeze(1).to(dtype=x.dtype, device=x.device)
+    x_even = x[..., 0::2]
+    x_odd = x[..., 1::2]
+    out = torch.empty_like(x)
+    out[..., 0::2] = x_even * c - x_odd * s
+    out[..., 1::2] = x_even * s + x_odd * c
+    return out
+
+
 class CausalSelfAttention(nn.Module):
     """Causal multi-head self-attention with SDPA and 2D-RoPE."""
 
@@ -117,17 +131,26 @@ class CausalSelfAttention(nn.Module):
         self.head_dim = config.n_embd // config.n_head
         self.dropout = config.dropout
         self.rope_theta = config.rope_theta
+        self.pos_emb_type = config.pos_emb_type
+        # 1D rotates head_dim/2 pairs; 2D rotates head_dim/4 pairs per subspace
+        self.rope_dim = 2 * self.head_dim if self.pos_emb_type == "rope_1d" else self.head_dim
 
         self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
         self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
         self.resid_dropout = nn.Dropout(config.dropout)
 
         cos, sin = precompute_2d_rope_freqs(
-            self.head_dim, max_coords=config.max_coords, theta_base=config.rope_theta
+            self.rope_dim, max_coords=config.max_coords, theta_base=config.rope_theta
         )
         self.register_buffer("rope_cos", cos, persistent=False)
         self.register_buffer("rope_sin", sin, persistent=False)
 
+    def rotate(self, x: torch.Tensor, place_coords: torch.Tensor, step_coords: torch.Tensor) -> torch.Tensor:
+        """Applies the configured RoPE variant (step_coords is ignored for rope_1d)."""
+        if self.pos_emb_type == "rope_1d":
+            return apply_1d_rope(x, place_coords, self.rope_cos, self.rope_sin)
+        return apply_2d_rope(x, place_coords, step_coords, self.rope_cos, self.rope_sin)
+
     def forward(
         self,
         x: torch.Tensor,
@@ -146,12 +169,12 @@ class CausalSelfAttention(nn.Module):
         max_coord = max(int(place_coords.max().item()), int(step_coords.max().item())) + 1
         if max_coord > self.rope_cos.size(0):
             new_max = max(max_coord * 2, 4096)
-            cos, sin = precompute_2d_rope_freqs(self.head_dim, max_coords=new_max, theta_base=self.rope_theta)
+            cos, sin = precompute_2d_rope_freqs(self.rope_dim, max_coords=new_max, theta_base=self.rope_theta)
             self.rope_cos = cos.to(x.device)
             self.rope_sin = sin.to(x.device)
 
-        q = apply_2d_rope(q, place_coords, step_coords, self.rope_cos, self.rope_sin)
-        k = apply_2d_rope(k, place_coords, step_coords, self.rope_cos, self.rope_sin)
+        q = self.rotate(q, place_coords, step_coords)
+        k = self.rotate(k, place_coords, step_coords)
 
         dropout_p = self.dropout if self.training else 0.0
         y = F.scaled_dot_product_attention(
@@ -202,6 +225,7 @@ class ScaledGPT(nn.Module):
 
     def __init__(self, config: ScaledGPTConfig, tokenizer: Optional[ScaledTokenizer] = None):
         super().__init__()
+        assert config.pos_emb_type in ("rope_2d", "rope_1d"), f"Unsupported pos_emb_type: {config.pos_emb_type}"
         self.config = config
         self.tokenizer = tokenizer or ScaledTokenizer()
 
@@ -237,6 +261,14 @@ class ScaledGPT(nn.Module):
             n_params -= self.transformer.wte.weight.numel()
         return n_params
 
+    def compute_coords(self, idx: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Returns (place, step) coordinates: 2D mapper for rope_2d, (0..T-1, zeros) for rope_1d."""
+        if self.config.pos_emb_type == "rope_1d":
+            B, T = idx.size()
+            pos = torch.arange(T, dtype=torch.long, device=idx.device).unsqueeze(0).expand(B, T)
+            return pos, torch.zeros_like(pos)
+        return self.compute_2d_coords(idx)
+
     def compute_2d_coords(self, idx: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         Unified 2D coordinate mapper:
@@ -360,11 +392,12 @@ class ScaledGPT(nn.Module):
         device = idx.device
 
         if place_coords is None or step_coords is None:
-            place_coords, step_coords = self.compute_2d_coords(idx)
+            place_coords, step_coords = self.compute_coords(idx)
 
         # Shift invariance data augmentation during training:
-        # Strictly applies selective shift masking (place_coords > 0).long()
-        if self.training and self.config.k_rand_shift > 0:
+        # Strictly applies selective shift masking (place_coords > 0).long().
+        # Skipped for rope_1d: a uniform offset leaves every relative distance, hence every RoPE score, unchanged.
+        if self.training and self.config.k_rand_shift > 0 and self.config.pos_emb_type == "rope_2d":
             shifts = torch.randint(0, self.config.k_rand_shift + 1, (B, 1), device=device)
             shift_mask = (place_coords > 0).long()
             place_coords = place_coords + shifts * shift_mask