By Daniil Yurshevich, Nikita Khomich

TLDR: we train LSTM model on algorithmic task of modulo addition and observe grokking. We fully reverse engeneer the algorithm learned and propose a way simpler equivalent version of the model that groks as well.

Reproducibility statement: all the code is available at the repo.

Introduction

This post is related to Neel Nanda's post and a detailed description of what grokking is can be found there. The short summary is that grokking is the phenomenon when model when being trained on an algorithmic task of relatively small size initially memorizes the trining set and then suddenly generalizes to the data it hasn't seen before. In our work we train a version of LSTM on modulo addition and observe grokking.

Model architecture, experiment settings and naming conventions

We train a version of LSTM with ReLU activation instead of $t a n h$ for getting hidden state from cell state $c_{t}$ . We use a linear layer without bias to get logits from $h_{2}$ . We use a slight trick which is using two linear layers without activation inbetween instad of one which makes grokking easier. We also use hidden dimension of $N = p$ and one-hot encoded embeddings. The exact model arcitecture as well as naming is (very similar to Chris Olah's amazing post)

We use similar parameters to Neel's paper meaning extremely large weight decay and particular set of betas for Adam found with grid search. We train the model on a 30% of all pairs $(a, b) m o d 113$ and observe grokking:

This is a typical and well studied behaviour for transformer, however we are not aware of any examples of reverse enginnered problem with observed grokking during training.

Exact formulas for what LSTM does

This section is only relevant to the methodology we used to reverse engineer the model, not understanding algorithm itself as well as clarifying all the architectural details used, so feel free to skip this section.

Below are the formulas describing all intermediate activations and operations performed by the given model architecture. The model predicts $a + b mod 113$ from inputs $a, b \in {0, \dots, 112} .$

Model Parameters and Dimensions

N = 113 (the modulo and also number of classes)

$vocab_size = N$

$hidden_size = 113$

The input is a pair $(a, b)$ with $a, b \in {0, \dots, N - 1} .$

Input Encoding

We have an input sequence of lengt...

Posts

Wikitag Contributions

Comments