latent_backtrack / scripts /make_diag_split_configs.py
Avra98's picture
Add training code (same as GitHub reasoning-by-superposition-latent)
8f46582 verified
Raw
History Blame Contribute Delete
3.1 kB
#!/usr/bin/env python3
"""Write L15 diag configs with SEPARATE promote vs backtrack gates.
Naming: diag_L15_prom{F|CE}{thr}_bt{F|CE|NONE}{thr}
promF095 = promote on hop-accuracy (frontier) >= 0.95
btCE090 = backtrack on both-arms score (ce_score) < 0.90
btNONE = no backtracking
"""
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1] / "args"
TEMPLATE = """\
# L15 DIAG — SEPARATE gates
# PROMOTE (stage i -> i+1): {promote_metric} >= {promote_threshold}
# BACKTRACK (retrain earlier): {bt_desc}
project: coconut
save_path: ckpts
name: "{name}"
only_eval: False
coconut: True
cot: False
no_thoughts: False
no_cot: False
c_thought: 1
max_latent_stage: 15
pad_latent_to_max: True
accuracy_staging: True
init_stage: 0
promote_metric: {promote_metric}
promote_threshold: {promote_threshold}
promote_on_current_only: {promote_on_current_only}
epochs_per_stage: 25
backtrack: {backtrack}
backtrack_metric: {backtrack_metric}
backtrack_detect_threshold: {backtrack_detect_threshold}
remember_rate: 0.3
revert_staging: False
eval_every: 10
log_every: 5
perhop_val_samples: 256
perhop_train_samples: 64
eval_print_full: False
backprop_depth: null
train_size: 50000
save_only_improve: False
save_every: 200
uniform_prob: 0.1
model_id: configs/symbol-2layer-8head-768dim-L20.json
load_model_path: ckpts/star-coconut-L15-bfs-stage0-warm/checkpoint_99
seed: 0
resume: 0
bf16: True
train_path: data/star_2arm_L15_train_fo_bfs.json
val_path: data/star_2arm_L15_valid_fo_bfs.json
reset_optimizer: False
batch_size_training: 128
debug: False
gradient_accumulation_steps: 1
num_epochs: 3000
lr: !!float "1e-4"
grad_clip: !!float "1.0"
warmup_steps: 200
weight_decay: 0.01
bfs_variant: True
"""
# (basename, promote_metric, promote_thr, bt_on, bt_metric, bt_thr, prom_cur_only)
ARMS = [
("diag_L15_promF095_btCE050", "frontier", 0.95, True, "ce_score", 0.50, False),
("diag_L15_promF095_btCE090", "frontier", 0.95, True, "ce_score", 0.90, False),
("diag_L15_promF095_btCE095", "frontier", 0.95, True, "ce_score", 0.95, False),
("diag_L15_promF099_btCE090", "frontier", 0.99, True, "ce_score", 0.90, False),
("diag_L15_promF095_btF095", "frontier", 0.95, True, "frontier", 0.95, False),
("diag_L15_promCE090_btCE090", "ce_score", 0.90, True, "ce_score", 0.90, False),
("diag_L15_promF095_btNONE", "frontier", 0.95, False, "frontier", 0.95, True),
("diag_L15_promF099_btCE050", "frontier", 0.99, True, "ce_score", 0.50, False),
]
def main():
for name, pm, pt, bt, bm, bt_thr, cur_only in ARMS:
bt_desc = "OFF" if not bt else f"{bm} < {bt_thr} triggers retrain"
text = TEMPLATE.format(
name=name,
promote_metric=pm,
promote_threshold=pt,
promote_on_current_only=str(cur_only),
backtrack=str(bt),
backtrack_metric=bm,
backtrack_detect_threshold=bt_thr,
bt_desc=bt_desc,
)
path = ROOT / f"{name}.yaml"
path.write_text(text)
print("wrote", path.name)
if __name__ == "__main__":
main()