Re-implementation of the paper ‘Grokking: Generalization beyond overfitting on small algorithmic datasets’
Paper
Original paper can be found here
Datasets
All datasets from the original paper’s appendix are supported.
Hyperparameters
The default hyperparameters are from the paper, but can be adjusted via the command line when running train.py
Running experiments
To run with default settings, simply run python train.py.
The first time you train on any dataset you have to specify --force_data.
Arguments:
optimizer args
- “–lr”, type=float, default=1e-3
- “–weight_decay”, type=float, default=1
- “–beta1”, type=float, default=0.9
- “–beta2”, type=float, default=0.98
model args
- “–num_heads”, type=int, default=4
- “–layers”, type=int, default=2
- “–width”, type=int, default=128
data args
- “–data_name”, type=str, default=”perm”, choices=[
- “perm_xy”, # permutation composition x * y
- “perm_xyx1”, # permutation composition x * y * x^-1
- “perm_xyx”, # permutation composition x * y * x
- “plus”, # x + y
- “minus”, # x – y
- “div”, # x / y
- “div_odd”, # x / y if y is odd else x – y
- “x2y2”, # x^2 + y^2
- “x2xyy2”, # x^2 + y^2 + xy
- “x2xyy2x”, # x^2 + y^2 + xy + x
- “x3xy”, # x^3 + y
- “x3xy2y” # x^3 + xy^2 + y
]
- “–num_elements”, type=int, default=5 (choose 5 for permutation data, 97 for arithmetic data)
- “–data_dir”, type=str, default=”./data”
- “–force_data”, action=”store_true”, help=”Whether to force dataset creation.”
training args
- “–batch_size”, type=int, default=512
- “–steps”, type=int, default=10**5
- “–train_ratio”, type=float, default=0.5
- “–seed”, type=int, default=42
- “–verbose”, action=”store_true”
- “–log_freq”, type=int, default=10
- “–num_workers”, type=int, default=4
- “–disable_logging”, action=”store_true”, help=”Whether to use wandb logging”
- “–checkpoints”, type=int, default=None, nargs=”*”, help=”List of number of steps after which to save model.”
