Codebase for Toward Understanding In-context vs. In-weight Learning (ICLR 2025)

Link: https://arxiv.org/abs/2410.23042

Python 3.10

pip install tensorflow-cpu tensorflow-datasets
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
pip install jax==0.4.30 # Or pip install -U "jax[cuda12]==0.4.30"
pip install flax==0.8.5 orbax-checkpoint==0.4.3
pip install chex optax dill gymnasium scikit-learn matplotlib seaborn
pip install prefetch_generator
pip install numpy==1.26.4

python src/main.py --config_path=experiments/config/simple_icl_prob_1.0.json

All experiments can be run using files in the local_utils directory. Important: Modify constants.py to setup the correct credentials and paths.

1. Generate the scripts using generate_train.py:

python generate_train.py

For each experiment EXP_NAME defined in configs.py, generate_train.py will generate a corresponding experiment file that is located in CONFIG_DIR (e.g. <CONFIG_DIR>/<EXP_NAME>.dat), and a corresponding bash script ./sbatch_scripts/run_all-<EXP_NAME>.sh. generate_train.py will also generate a bash script sbatch_all_train.sh that will kick off all experiments.

1. Kick off experiments using sbatch_all_train.sh:

chmod +x sbatch_all_train.sh
./sbatch_all_train.sh

You can use fine_tuning_tutorial.ipynb to fine-tune a Gemma model using only 10 fine-tuning prompts.

To visualize the results, use files in the plot_utils directory:

1. Run convert_to_pd.ipynb to convert runs into Panda Dataframe, saved in feather format.
2. Run plot_experiment.ipynb to visualize particular feather files

You may use theoretical_plots.ipynb to obtain the theoretical bound plots (using a more "general" estimator rather than the KT estimator).

If you reuse this code from this repository, please cite our paper Toward Understanding In-context vs. In-weight Learning:

@inproceedings{chan2025towardicl,
    title={Toward Understanding In-context vs. In-weight Learning},
    author={Chan, Bryan and Chen, Xinyi and Gy{\"o}rgy, Andr{\'a}s and Schuurmans, Dale},
    year={2025},
    booktitle={International Conference on Learning Representations (ICLR)},
}