# 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.”