In-context Learning task

[ldqvinh]

Hello, Thanks for your work. I attempted the in-context learning training command from the experiment details, but encountered a 'loss is NaN' error. Could you share the command you used? Appreciate it.

python src/run.py \
     --dataset "induction" \
     --vocab_size 10 \
     --dataset_size 20000 \
     --min_length 1 \
     --max_length 10 \
     --n_epochs 250 \
     --batch_size 512 \
     --lr "5e-2" \
     --n_layers 2 \
     --n_heads_cat 1 \
     --n_heads_num 0 \
     --n_cat_mlps 1 \
     --n_num_mlps 0 \
     --one_hot_embed \
     --count_only \
     --seed 0 \
     --save \
     --save_code \
     --output_dir "output/induction"

[danfriedman0]

Thanks for taking an interest in the code! I'm not immediately sure what the issue could be, but some things to try are:

1. Set --min_length 10 and --max_length 10
2. Set the --autoregressive flag.
3. Set --n_cat_mlps 0 (no MLPs)
4. Set --n_epochs = 500.

Note that you likely need to try a number of random seeds to get a model that successfully learns the task. To save time, we also used a "patience" of 25 (this is a possible argument to the run_training function, although you would need to modify src/run.py to make it a command-line flag).

Could you also share any more details about what you observe? Do you get "loss is NaN" right away, or only after some training?

[Wangcheng-Xu]

Hi @danfriedman0,

I also had issue replicating the induction experiment. The command was as suggested above, which is also copied below. I used a modified file "experiment_run_n.py" that iterates through seed when a training running out of the patience with an additional "patience" argument. The training seems to return a constant loss=5.81e+29 from seed 0 all the way to 100. By the way, some other experiments seemed to work, such as "sort", "reverse", etc.

CUDA_VISIBLE_DEVICES=0 python experiment_run_n.py \ --dataset "induction" \ --vocab_size 10 \ --dataset_size 20000 \ --min_length 10 \ --max_length 10 \ --n_epochs 500 \ --batch_size 512 \ --patience 25 \ --lr "5e-2" \ --n_layers 2 \ --n_heads_cat 1 \ --n_heads_num 0 \ --n_cat_mlps 0 \ --n_num_mlps 0 \ --one_hot_embed \ --count_only \ --autoregressive \ --seed 0 \ --save \ --save_code \ --output_dir "output/induction"

Also, it would be great if you can share the configuration for replicating all experiments in the paper, like that for "sort" and "conll_ner" in the README.md. Thanks!

[danfriedman0]

Hi all, sorry for the trouble, and thanks for the additional detail.

I think I found the main problem: you need to set --unembed_mask 0. This flag is set to 1 by default, which prevents the model from predicting pad or unk as the output token, but the unk token is a valid prediction for this task. I have uploaded a script with a command that works for me (on around 20% of seeds).

@Wangcheng-Xu : The scripts directory contains configurations used for the other experiments in the paper. Please let me know if you have any more questions.

[Wangcheng-Xu]

Thank you! I have tested the fixed configuration for the induction task, which works for me.

[ldqvinh]

Thank you to everyone involved for identifying and resolving the issue. The updated configuration for the induction task is now functioning perfectly on my end as well.