Controlling Grokking with Nonlinearity and Data Symmetry
This work addresses grokking in neural networks for researchers, offering incremental insights into controlling generalization through architectural modifications.
The paper tackles controlling grokking behavior in neural networks for modular arithmetic by adjusting activation functions, depth, and width, resulting in patterns that can factor nonprime moduli and metrics linking generalization to weight entropy.
This paper demonstrates that grokking behavior in modular arithmetic with a modulus P in a neural network can be controlled by modifying the profile of the activation function as well as the depth and width of the model. Plotting the even PCA projections of the weights of the last NN layer against their odd projections further yields patterns which become significantly more uniform when the nonlinearity is increased by incrementing the number of layers. These patterns can be employed to factor P when P is nonprime. Finally, a metric for the generalization ability of the network is inferred from the entropy of the layer weights while the degree of nonlinearity is related to correlations between the local entropy of the weights of the neurons in the final layer.