Upload add_sub_baseline_10K_2L1H128d

add_sub_baseline_10K_2L1H128d/config.json

@@ -0,0 +1,37 @@
+{
+  "architectures": [
+    "SorlModelWrapper"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": null,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 512,
+  "layer_types": [
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 128,
+  "max_window_layers": 28,
+  "model_type": "qwen3",
+  "num_attention_heads": 1,
+  "num_hidden_layers": 2,
+  "num_key_value_heads": 1,
+  "pad_token_id": null,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.5.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151645
+}

add_sub_baseline_10K_2L1H128d/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.5.0",
+  "use_cache": true
+}

add_sub_baseline_10K_2L1H128d/metrics.json

@@ -0,0 +1,831 @@
+{
+  "history": {
+    "step": [
+      50,
+      100,
+      150,
+      200,
+      250,
+      300,
+      350,
+      400,
+      450,
+      500,
+      550,
+      600,
+      650,
+      700,
+      750,
+      800,
+      850,
+      900,
+      950,
+      1000,
+      1050,
+      1100,
+      1150,
+      1200,
+      1250,
+      1300,
+      1350,
+      1400,
+      1450,
+      1500,
+      1550,
+      1600,
+      1650,
+      1700,
+      1750,
+      1800,
+      1850,
+      1900,
+      1950,
+      2000,
+      2050,
+      2100,
+      2150,
+      2200,
+      2250,
+      2300,
+      2350,
+      2400,
+      2450,
+      2500,
+      2550,
+      2600,
+      2650,
+      2700,
+      2750,
+      2800,
+      2850,
+      2900,
+      2950,
+      3000,
+      3050,
+      3100
+    ],
+    "loss": [
+      11.632269859313965,
+      10.746880531311035,
+      10.444467544555664,
+      10.006251335144043,
+      9.801725387573242,
+      9.507078170776367,
+      9.177903175354004,
+      8.873117446899414,
+      8.627289772033691,
+      8.425443649291992,
+      8.176294326782227,
+      7.824686527252197,
+      7.672322750091553,
+      7.39719820022583,
+      7.087851524353027,
+      6.878063678741455,
+      6.657628059387207,
+      6.365849018096924,
+      6.228781223297119,
+      5.940927982330322,
+      5.686161994934082,
+      5.552325248718262,
+      5.286587715148926,
+      5.14711856842041,
+      4.970963001251221,
+      4.826735019683838,
+      4.654284954071045,
+      4.47059965133667,
+      4.298427581787109,
+      4.127253532409668,
+      4.005905628204346,
+      3.8805015087127686,
+      3.801790475845337,
+      3.6226108074188232,
+      3.5523719787597656,
+      3.5552258491516113,
+      3.372868061065674,
+      3.3437774181365967,
+      3.2328782081604004,
+      3.155165910720825,
+      3.211038112640381,
+      3.113257884979248,
+      2.977780818939209,
+      3.0747323036193848,
+      2.98803448677063,
+      2.9421255588531494,
+      2.9898440837860107,
+      2.9196650981903076,
+      2.8956546783447266,
+      2.9016549587249756,
+      2.8742430210113525,
+      2.872373342514038,
+      2.8761067390441895,
+      2.807882070541382,
+      2.8126957416534424,
+      2.813640594482422,
+      2.8270370960235596,
+      2.8576278686523438,
+      2.861673593521118,
+      2.809460401535034,
+      2.736629009246826,
+      2.832703113555908
+    ],
+    "base_loss": [
+      11.632269859313965,
+      10.746880531311035,
+      10.444467544555664,
+      10.006251335144043,
+      9.801725387573242,
+      9.507078170776367,
+      9.177903175354004,
+      8.873117446899414,
+      8.627289772033691,
+      8.425443649291992,
+      8.176294326782227,
+      7.824686527252197,
+      7.672322750091553,
+      7.39719820022583,
+      7.087851524353027,
+      6.878063678741455,
+      6.657628059387207,
+      6.365849018096924,
+      6.228781223297119,
+      5.940927982330322,
+      5.686161994934082,
+      5.552325248718262,
+      5.286587715148926,
+      5.14711856842041,
+      4.970963001251221,
+      4.826735019683838,
+      4.654284954071045,
+      4.47059965133667,
+      4.298427581787109,
+      4.127253532409668,
+      4.005905628204346,
+      3.8805015087127686,
+      3.801790475845337,
+      3.6226108074188232,
+      3.5523719787597656,
+      3.5552258491516113,
+      3.372868061065674,
+      3.3437774181365967,
+      3.2328782081604004,
+      3.155165910720825,
+      3.211038112640381,
+      3.113257884979248,
+      2.977780818939209,
+      3.0747323036193848,
+      2.98803448677063,
+      2.9421255588531494,
+      2.9898440837860107,
+      2.9196650981903076,
+      2.8956546783447266,
+      2.9016549587249756,
+      2.8742430210113525,
+      2.872373342514038,
+      2.8761067390441895,
+      2.807882070541382,
+      2.8126957416534424,
+      2.813640594482422,
+      2.8270370960235596,
+      2.8576278686523438,
+      2.861673593521118,
+      2.809460401535034,
+      2.736629009246826,
+      2.832703113555908
+    ],
+    "lr": [
+      9.800000000000001e-06,
+      1.98e-05,
+      1.9987181950611177e-05,
+      1.9947710543876128e-05,
+      1.9881685800704465e-05,
+      1.978928396036359e-05,
+      1.9670751670274858e-05,
+      1.9526405327639747e-05,
+      1.9356630234883628e-05,
+      1.916187957117136e-05,
+      1.8942673182740125e-05,
+      1.8699596195278423e-05,
+      1.8433297452055233e-05,
+      1.814448778196835e-05,
+      1.7833938102135036e-05,
+      1.7502477360089763e-05,
+      1.7150990321081803e-05,
+      1.6780415206379127e-05,
+      1.6391741188882608e-05,
+      1.598600575273548e-05,
+      1.556429192397602e-05,
+      1.5127725379625657e-05,
+      1.4677471442929161e-05,
+      1.4214731972767564e-05,
+      1.3740742155546771e-05,
+      1.3256767208125341e-05,
+      1.2764099000582231e-05,
+      1.2264052607839339e-05,
+      1.1757962799343548e-05,
+      1.124718047617836e-05,
+      1.0733069065115504e-05,
+      1.021700087923181e-05,
+      9.700353454805953e-06,
+      9.184505874272968e-06,
+      8.67083508505161e-06,
+      8.160712224070698e-06,
+      7.655498957805332e-06,
+      7.156543847592537e-06,
+      6.6651787499283185e-06,
+      6.182715261354808e-06,
+      5.710441217427128e-06,
+      5.249617255105255e-06,
+      4.801473447746826e-06,
+      4.3672060216831465e-06,
+      3.947974163142735e-06,
+      3.5448969240456945e-06,
+      3.1590502349282503e-06,
+      2.791464032970812e-06,
+      2.4431195127956985e-06,
+      2.1149465073729436e-06,
+      1.8078210060253153e-06,
+      1.522562816157701e-06,
+      1.2599333749523968e-06,
+      1.0206337168715263e-06,
+      8.05302602391903e-07,
+      6.145148129672918e-07,
+      4.4877961676932255e-07,
+      3.0853940930242963e-07,
+      1.941685325214182e-07,
+      1.05972275603764e-07,
+      4.4186060043894804e-08,
+      8.97481124466415e-09
+    ],
+    "eval_step": [
+      156,
+      312,
+      468,
+      624,
+      780,
+      936,
+      1092,
+      1248,
+      1404,
+      1560,
+      1716,
+      1872,
+      2028,
+      2184,
+      2340,
+      2496,
+      2652,
+      2808,
+      2964,
+      3120
+    ],
+    "eval_epoch": [
+      1,
+      2,
+      3,
+      4,
+      5,
+      6,
+      7,
+      8,
+      9,
+      10,
+      11,
+      12,
+      13,
+      14,
+      15,
+      16,
+      17,
+      18,
+      19,
+      20
+    ],
+    "eval_accuracy": [
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0
+    ]
+  },
+  "final_accuracy": 0.0,
+  "sft_eval": {
+    "config": {
+      "ops": "add_sub",
+      "K": null,
+      "mode": "sft",
+      "n_digits": 6,
+      "n_per_split": 100
+    },
+    "splits": {
+      "add_S0": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "SA": {
+            "accuracy": 0.17851239669421487,
+            "count": 605
+          },
+          "SS": {
+            "accuracy": 0.0,
+            "count": 95
+          }
+        }
+      },
+      "add_S1": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "SA": {
+            "accuracy": 0.24509803921568626,
+            "count": 204
+          },
+          "SC": {
+            "accuracy": 0.22485207100591717,
+            "count": 169
+          },
+          "SS": {
+            "accuracy": 0.0,
+            "count": 31
+          },
+          "UC": {
+            "accuracy": 0.006756756756756757,
+            "count": 296
+          }
+        }
+      },
+      "add_S2": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "SA": {
+            "accuracy": 0.2331288343558282,
+            "count": 163
+          },
+          "SC": {
+            "accuracy": 0.2,
+            "count": 130
+          },
+          "SS": {
+            "accuracy": 0.0,
+            "count": 87
+          },
+          "UC": {
+            "accuracy": 0.0,
+            "count": 203
+          },
+          "US": {
+            "accuracy": 1.0,
+            "count": 117
+          }
+        }
+      },
+      "add_S3": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "SA": {
+            "accuracy": 0.2727272727272727,
+            "count": 121
+          },
+          "SC": {
+            "accuracy": 0.1652892561983471,
+            "count": 121
+          },
+          "SS": {
+            "accuracy": 0.0,
+            "count": 49
+          },
+          "UC": {
+            "accuracy": 0.0,
+            "count": 186
+          },
+          "US": {
+            "accuracy": 1.0,
+            "count": 223
+          }
+        }
+      },
+      "add_S4": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "SA": {
+            "accuracy": 0.2403846153846154,
+            "count": 104
+          },
+          "SC": {
+            "accuracy": 0.22641509433962265,
+            "count": 106
+          },
+          "SS": {
+            "accuracy": 0.0,
+            "count": 23
+          },
+          "UC": {
+            "accuracy": 0.0,
+            "count": 160
+          },
+          "US": {
+            "accuracy": 1.0,
+            "count": 307
+          }
+        }
+      },
+      "add_S5": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "SA": {
+            "accuracy": 0.5,
+            "count": 100
+          },
+          "SC": {
+            "accuracy": 0.28,
+            "count": 100
+          },
+          "UC": {
+            "accuracy": 0.0,
+            "count": 100
+          },
+          "US": {
+            "accuracy": 1.0,
+            "count": 400
+          }
+        }
+      },
+      "add_S6": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "SC": {
+            "accuracy": 0.21,
+            "count": 100
+          },
+          "UC": {
+            "accuracy": 0.0,
+            "count": 100
+          },
+          "US": {
+            "accuracy": 1.0,
+            "count": 500
+          }
+        }
+      },
+      "add_random": {
+        "full_accuracy": 0.0,
+        "n_examples": 200,
+        "per_subtask": {
+          "SA": {
+            "accuracy": 0.2371364653243848,
+            "count": 447
+          },
+          "SC": {
+            "accuracy": 0.2125,
+            "count": 320
+          },
+          "SS": {
+            "accuracy": 0.0,
+            "count": 56
+          },
+          "UC": {
+            "accuracy": 0.007561436672967864,
+            "count": 529
+          },
+          "US": {
+            "accuracy": 1.0,
+            "count": 48
+          }
+        }
+      },
+      "add_C3": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "SA": {
+            "accuracy": 0.31666666666666665,
+            "count": 300
+          },
+          "SC": {
+            "accuracy": 0.26,
+            "count": 100
+          },
+          "UC": {
+            "accuracy": 0.0051813471502590676,
+            "count": 193
+          },
+          "US": {
+            "accuracy": 1.0,
+            "count": 107
+          }
+        }
+      },
+      "add_C4": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "SA": {
+            "accuracy": 0.415,
+            "count": 200
+          },
+          "SC": {
+            "accuracy": 0.21,
+            "count": 100
+          },
+          "UC": {
+            "accuracy": 0.0,
+            "count": 256
+          },
+          "US": {
+            "accuracy": 1.0,
+            "count": 144
+          }
+        }
+      },
+      "add_C5": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "SA": {
+            "accuracy": 0.64,
+            "count": 100
+          },
+          "SC": {
+            "accuracy": 0.23,
+            "count": 100
+          },
+          "UC": {
+            "accuracy": 0.0,
+            "count": 306
+          },
+          "US": {
+            "accuracy": 1.0,
+            "count": 194
+          }
+        }
+      },
+      "add_C6": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "SC": {
+            "accuracy": 0.24,
+            "count": 100
+          },
+          "UC": {
+            "accuracy": 0.0,
+            "count": 366
+          },
+          "US": {
+            "accuracy": 1.0,
+            "count": 234
+          }
+        }
+      },
+      "sub_M0": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "MD": {
+            "accuracy": 0.20465890183028287,
+            "count": 601
+          },
+          "ME": {
+            "accuracy": 1.0,
+            "count": 99
+          }
+        }
+      },
+      "sub_M1": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "MD": {
+            "accuracy": 0.3835125448028674,
+            "count": 279
+          },
+          "MB": {
+            "accuracy": 0.0,
+            "count": 145
+          },
+          "ME": {
+            "accuracy": 1.0,
+            "count": 24
+          },
+          "UB": {
+            "accuracy": 0.09523809523809523,
+            "count": 252
+          }
+        }
+      },
+      "sub_M2": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "MD": {
+            "accuracy": 0.6150234741784038,
+            "count": 213
+          },
+          "MB": {
+            "accuracy": 0.0,
+            "count": 113
+          },
+          "ME": {
+            "accuracy": 1.0,
+            "count": 85
+          },
+          "UB": {
+            "accuracy": 0.16574585635359115,
+            "count": 181
+          },
+          "UD": {
+            "accuracy": 0.0,
+            "count": 108
+          }
+        }
+      },
+      "sub_M3": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "MD": {
+            "accuracy": 0.7597765363128491,
+            "count": 179
+          },
+          "MB": {
+            "accuracy": 0.0,
+            "count": 103
+          },
+          "ME": {
+            "accuracy": 1.0,
+            "count": 56
+          },
+          "UB": {
+            "accuracy": 0.12080536912751678,
+            "count": 149
+          },
+          "UD": {
+            "accuracy": 0.0,
+            "count": 213
+          }
+        }
+      },
+      "sub_M4": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "MD": {
+            "accuracy": 0.5,
+            "count": 200
+          },
+          "MB": {
+            "accuracy": 0.0,
+            "count": 100
+          },
+          "UB": {
+            "accuracy": 0.3,
+            "count": 100
+          },
+          "UD": {
+            "accuracy": 0.0,
+            "count": 300
+          }
+        }
+      },
+      "sub_M5": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "MD": {
+            "accuracy": 1.0,
+            "count": 100
+          },
+          "MB": {
+            "accuracy": 0.0,
+            "count": 100
+          },
+          "UB": {
+            "accuracy": 0.31,
+            "count": 100
+          },
+          "UD": {
+            "accuracy": 0.0,
+            "count": 400
+          }
+        }
+      },
+      "sub_random": {
+        "full_accuracy": 0.0,
+        "n_examples": 200,
+        "per_subtask": {
+          "MD": {
+            "accuracy": 0.3616666666666667,
+            "count": 600
+          },
+          "MB": {
+            "accuracy": 0.0,
+            "count": 267
+          },
+          "ME": {
+            "accuracy": 1.0,
+            "count": 53
+          },
+          "UB": {
+            "accuracy": 0.12072892938496584,
+            "count": 439
+          },
+          "UD": {
+            "accuracy": 0.0,
+            "count": 41
+          }
+        }
+      },
+      "sub_B3": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "MD": {
+            "accuracy": 0.3333333333333333,
+            "count": 300
+          },
+          "MB": {
+            "accuracy": 0.0,
+            "count": 100
+          },
+          "UB": {
+            "accuracy": 0.17766497461928935,
+            "count": 197
+          },
+          "UD": {
+            "accuracy": 0.0,
+            "count": 103
+          }
+        }
+      },
+      "sub_B4": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "MD": {
+            "accuracy": 0.5,
+            "count": 200
+          },
+          "MB": {
+            "accuracy": 0.0,
+            "count": 100
+          },
+          "UB": {
+            "accuracy": 0.145748987854251,
+            "count": 247
+          },
+          "UD": {
+            "accuracy": 0.0,
+            "count": 153
+          }
+        }
+      },
+      "sub_B5": {
+        "full_accuracy": 0.0,
+        "n_examples": 100,
+        "per_subtask": {
+          "MD": {
+            "accuracy": 1.0,
+            "count": 100
+          },
+          "MB": {
+            "accuracy": 0.0,
+            "count": 100
+          },
+          "UB": {
+            "accuracy": 0.11073825503355705,
+            "count": 298
+          },
+          "UD": {
+            "accuracy": 0.0,
+            "count": 202
+          }
+        }
+      }
+    },
+    "summary": {
+      "overall_accuracy": 0.0,
+      "total_examples": 2400,
+      "n_splits": 22
+    }
+  }
+}

add_sub_baseline_10K_2L1H128d/model.safetensors

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

add_sub_baseline_10K_2L1H128d/train_config.json

@@ -0,0 +1,84 @@
+{
+  "num_rollouts": 4,
+  "K": 4,
+  "max_iterations": 2,
+  "memory_span_abs": 1792,
+  "memory_span_traj": 1792,
+  "temperature": 1.0,
+  "ar_search": false,
+  "response_only_abs": false,
+  "alpha_info_gain": 10.0,
+  "alpha_abs": 0.1,
+  "alpha_soft_zipf": 1.0,
+  "alpha_ortho": 0.0,
+  "alpha_anchor": 0.0,
+  "alpha_jacobi": 0.0,
+  "decay": 0.8,
+  "target_vocab_util": 0.8,
+  "min_abs_ppl": 0.0,
+  "zipf_alpha": 1.0,
+  "lr": 2e-05,
+  "emb_lr_mult": 1.0,
+  "weight_decay": 0.01,
+  "warmup_steps": 100,
+  "cooldown_frac": 0.4,
+  "max_grad_norm": 1.0,
+  "vq_abs_pretrain_steps": 0,
+  "vq_abs_pretrain_lr": 0.001,
+  "vq_abs_pretrain_layer": -1,
+  "vq_abs_pretrain_batch_size": 256,
+  "vq_abs_pretrain_target_vectors": 20000,
+  "batch_size": 64,
+  "gradient_accumulation_steps": 1,
+  "num_epochs": 20,
+  "emb_warmup_steps": 0,
+  "log_every": 50,
+  "eval_every": 156,
+  "save_every": 999999,
+  "eval_samples": 100,
+  "output_dir": "ckpt/sweep/as_baseline_10K_2L1H128d",
+  "eval_K": 4,
+  "alpha_traj": 0.0,
+  "corrupt_method": "shuffle",
+  "corrupt_ratio": 0.3,
+  "alpha_contrastive": 1.0,
+  "gamma_contrastive": 0.5,
+  "alpha_masked_traj": 0.0,
+  "mask_nl_ratio": 0.3,
+  "mask_nl_mode": "fixed",
+  "mask_nl_fixed_id": 0,
+  "use_ste": true,
+  "n_inner": 1,
+  "random_K": null,
+  "strip_suffix": null,
+  "compress_prefix": null,
+  "random_mem_span": null,
+  "warmup_ratio": 0.03,
+  "beta2": 0.999,
+  "seed": 42,
+  "n_digits": 6,
+  "n_layer": 2,
+  "n_head": 1,
+  "n_embd": 128,
+  "ops": "add_sub",
+  "abs_vocab": 0,
+  "dataset_size": 10000,
+  "mode": "baseline",
+  "device": "cuda",
+  "push_to_hub": true,
+  "no_wandb": false,
+  "n_params": 39346560,
+  "run_name": "add_sub_baseline_10K_2L1H128d",
+  "git_commit": "7c8a203ce79a277a1e41a3ec0648cb73d5b2b760",
+  "timestamp": "2026-04-14T01:48:04.187677+00:00",
+  "tokenizer": "Qwen/Qwen3-0.6B",
+  "dataset_repo": "thoughtworks/arithmetic-sorl-data",
+  "dataset_config": "add_sub_6digit",
+  "model_repo": "thoughtworks/arithmetic-sorl",
+  "trainer_version": "sft",
+  "wandb_run_id": "mxgukq0x",
+  "wandb_url": "https://wandb.ai/nlp_and_interpretability/sorl-arithmetic/runs/mxgukq0x",
+  "final_accuracy": 0.0,
+  "sft_accuracy": 0.0,
+  "eval_method": "ArithmeticEvaluator"
+}