Robust generalization is a major challenge in deep learning, particularly when the number of trainable parameters is very large. In general, it is very difficult to know if the network has memorized a particular set of examples or understood the underlying rule (or both). Motivated by this challenge, we study an interpretable model where generalizing representations are understood analytically, and are easily distinguishable from the memorizing ones. Namely, we consider two-layer neural networks trained on modular arithmetic tasks where ($\xi \cdot 100\%$) of labels are corrupted (\emph{i.e.} some results of the modular operations in the training set are incorrect). We show that (i) it is possible for the network to memorize the corrupted labels \emph{and} achieve $100\%$ generalization at the same time; (ii) the memorizing neurons can be identified and pruned, lowering the accuracy on corrupted data and improving the accuracy on uncorrupted data; (iii) regularization methods such as weight decay, dropout and BatchNorm force the network to ignore the corrupted data during optimization, and achieve $100\%$ accuracy on the uncorrupted dataset; and (iv) the effect of these regularization methods is (``mechanistically'') interpretable: weight decay and dropout force all the neurons to learn generalizing representations, while BatchNorm de-amplifies the output of memorizing neurons and amplifies the output of the generalizing ones. Finally, we show that in the presence of regularization, the training dynamics involves two consecutive stages: first, the network undergoes the \emph{grokking} dynamics reaching high train \emph{and} test accuracy; second, it unlearns the memorizing representations, where train accuracy suddenly jumps from $100\%$ to $100 (1-\xi)\%$.
翻译:鲁棒泛化是深度学习中的一大挑战,尤其是在可训练参数数量极多的情况下。一般而言,很难判断网络是记忆了特定样本集,还是理解了底层规则(或两者兼有)。受这一挑战的启发,我们研究了一个可解释模型,在该模型中,泛化表征能被分析性地理解,且易于与记忆表征区分。具体而言,我们考虑在模算术任务上训练的两层神经网络,其中($\xi \cdot 100\%$)的标签被破坏(即训练集中部分模运算结果错误)。我们发现:(i) 网络有可能同时记忆被破坏标签并实现$100\%$泛化;(ii) 记忆性神经元可以被识别并剪枝,从而降低在破坏数据上的准确率,并提升在未破坏数据上的准确率;(iii) 权重衰减、Dropout和BatchNorm等正则化方法迫使网络在优化过程中忽略破坏数据,并在未破坏数据集上实现$100\%$准确率;(iv) 这些正则化方法的效果具有(“机制层面的”)可解释性:权重衰减和Dropout迫使所有神经元学习泛化表征,而BatchNorm则减弱记忆性神经元的输出并放大泛化性神经元的输出。最后,我们表明在存在正则化的情况下,训练动态包含两个连续阶段:首先,网络经历“顿悟”动态,达到高训练准确率和测试准确率;其次,它遗忘记忆表征,此时训练准确率突然从$100\%$下降到$100 (1-\xi)\%$。