Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks

This is an official implementation of the paper arXiv:2406.02550.

Requirements

• Tested Environment:
  • Python == 3.11.4
  • PyTorch >= 2.2
• Run pip install -r requirements.txt before using

Structure of this Repo

We offer separate README files in each folder.

• Tools: _src/
• Pre-training scripts: pre_training/
• Interpretability:
  • Fig 5: interpretability/...
  • Attention Map and PCA Analysis: interpretability/pca_analysis*