icl-dynamics/ih_paper_appendix_runs.sh at main · aadityasingh/icl-dynamics · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
#!/bin/bash

# This file reproduces all runs necessary to reproduce the figures in the main paper
# and some plots in the appendix

### Overarching params, copied from paper_runs.sh
SAVE_FOLDER="./ih_paper_reprod/main_paper_is5_ih3_pt2"
INIT_SEED=5
REST_HEADS="1:0 1:1 1:2 1:4 1:5 1:6 1:7"
STRONGEST_IH="1:3"
IHNAME=3

# Which other
ADDITIONAL_SEED_1=6
ADDITIONAL_SEED_2=7

# How long to run other init seeds for, to get checkpoints to graft from
MAIN_RUN_ITERS=1000000
GRAFT_CKPT="001000000" # Need 0 padding to have the correct filename to load

# We want to make sure clamping previous token attention pattern
# isn't sensitive to the specific head we clamped, so we try:
# clamping a different head's pattern
PT_CLAMP_SINGLE_OTHER="0:6"
# clamping multiple heads' patterns
PT_CLAMP_MULTI="0:1 0:2 0:5"
# clamping all heads' patterns
PT_CLAMP_ALL="0:0 0:1 0:2 0:3 0:4 0:5 0:6 0:7"

# How long to run clamp experiments for
CLAMP_ITERS=800000

# How long to run variable match experiments for
MATCH_ITERS=600000

### Figure 7
# prereq, additional seed runs, only need last ckpt
python main.py --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MAIN_RUN_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --run omniglot50_rl5_is$ADDITIONAL_SEED_1 --ckpt_every $MAIN_RUN_ITERS --init_seed $ADDITIONAL_SEED_1 --base_folder $SAVE_FOLDER --raw_name
python main.py --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MAIN_RUN_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --run omniglot50_rl5_is$ADDITIONAL_SEED_2 --ckpt_every $MAIN_RUN_ITERS --init_seed $ADDITIONAL_SEED_2 --base_folder $SAVE_FOLDER --raw_name

# Clamp runs with all 3 seeds (using graft functionality)
python main.py --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $CLAMP_ITERS --train_bs 32 --eval_every 1000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --run omniglot50_rl5_graft_l0_from_is$INIT_SEED --init_seed $INIT_SEED --opto_graft_in_model_ckpt $SAVE_FOLDER/omniglot50_rl5/checkpoints/$GRAFT_CKPT.eqx --opto_graft_in_model_cfg $SAVE_FOLDER/omniglot50_rl5/config.json --opto_graft_in_model_till_layer 0 --base_folder $SAVE_FOLDER --raw_name

python main.py --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $CLAMP_ITERS --train_bs 32 --eval_every 1000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --run omniglot50_rl5_graft_l0_from_is$ADDITIONAL_SEED_1 --init_seed $INIT_SEED --opto_graft_in_model_ckpt $SAVE_FOLDER/omniglot50_rl5_is$ADDITIONAL_SEED_1/checkpoints/$GRAFT_CKPT.eqx --opto_graft_in_model_cfg $SAVE_FOLDER/omniglot50_rl5_is$ADDITIONAL_SEED_1/config.json --opto_graft_in_model_till_layer 0 --base_folder $SAVE_FOLDER --raw_name

python main.py --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $CLAMP_ITERS --train_bs 32 --eval_every 1000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --run omniglot50_rl5_graft_l0_from_is$ADDITIONAL_SEED_2 --init_seed $INIT_SEED --opto_graft_in_model_ckpt $SAVE_FOLDER/omniglot50_rl5_is$ADDITIONAL_SEED_2/checkpoints/$GRAFT_CKPT.eqx --opto_graft_in_model_cfg $SAVE_FOLDER/omniglot50_rl5_is$ADDITIONAL_SEED_2/config.json --opto_graft_in_model_till_layer 0 --base_folder $SAVE_FOLDER --raw_name

### Figure 8, alternative PT attention clamps
python main.py --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $CLAMP_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --run omniglot50_rl5_perfect_prev_tok_attn_single_other --opto_prev_token_heads $PT_CLAMP_SINGLE_OTHER --opto_ablate_heads $REST_HEADS --init_seed $INIT_SEED --base_folder $SAVE_FOLDER --raw_name

python main.py --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $CLAMP_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --run omniglot50_rl5_perfect_prev_tok_attn_multi --opto_prev_token_heads $PT_CLAMP_MULTI --opto_ablate_heads $REST_HEADS --init_seed $INIT_SEED --base_folder $SAVE_FOLDER --raw_name

python main.py --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $CLAMP_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --run omniglot50_rl5_perfect_prev_tok_attn_all --opto_prev_token_heads $PT_CLAMP_ALL --opto_ablate_heads $REST_HEADS --init_seed $INIT_SEED --base_folder $SAVE_FOLDER --raw_name

### Figure 9, toy model run + config
python simple_model_solver.py --save ${SAVE_FOLDER}/figures/appendix_figure_9.pdf

### Figure 10 progress measures doesn't involve additional training runs, see jupyter notebook

### Figure 12 -- add this in jupyter notebook, don't rely on visualize runs

### Figure 13, 14 both require no additional runs

### Figure 15, lots of runs to sweep strengths
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 1 --run 100;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.95 --run 095;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.9 --run 090;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.85 --run 085;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.8 --run 080;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.75 --run 075;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.7 --run 070;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.65 --run 065;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.6 --run 060;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.55 --run 055;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.5 --run 050;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.4 --run 040;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.3 --run 030;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.2 --run 020;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.1 --run 010;
python main.py --base_folder $SAVE_FOLDER/ih"$IHNAME"_strength --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MATCH_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --opto_induction_heads $STRONGEST_IH --init_seed $INIT_SEED --raw_name --opto_induction_head_strength 0.0 --run 000;

### Figure 16 required no additional runs

### Figure 17 paired runs
arr=($REST_HEADS)

for ((i=0; i<${#arr[@]}; i++)); do
    unset arr[$i]
    python main.py --data_file omniglot_resnet18_randomized_order_s0.h5 --mixing_coeffs 1.0 --pt_burstiness 1 --train_context_len 2 --fs_relabel 5 --fs_relabel_split 8 2 0 --exemplar_split 1 4 0 --pe_names fsl_train fsl_val_rl fsl_train_valex fsl_test_class --pe_classes train train train test --pe_exemplars train train val train --pe_fs_relabel_scheme train val train train --pe_burstiness 1 1 1 1 --train_iters $MAIN_RUN_ITERS --train_bs 32 --eval_every 5000 --eval_iters 1000 --lr 0.00001 --d_model 64 --class_split 50 1473 100 --run omniglot50_rl5_only_l1h"$IHNAME"_and_iter$i --ckpt_every $MAIN_RUN_ITERS --opto_ablate_heads "${arr[@]}" --init_seed $INIT_SEED --base_folder $SAVE_FOLDER --raw_name
    arr=($REST_HEADS)
done