Do Mice Grok? Glimpses of Hidden Progress in Sensory Cortex

Thank you for the review, and we're sorry to hear you disliked our work. We've made substantial changes in response to your feedback. To clarify: this is not a work on "the intersection of foundation models and cognitive science" -- the paper is neither centrally about foundation models nor cognitive science. Instead, it argues that dynamics in existing neural data from genuine biological networks exhibit close similarities with a phenomenon hitherto seen and studied only in deep learning, presents new evidence for this case, and constructs and interprets simple machine learning models to flesh out a theoretical basis for the argument.

Writing Quality

We've made sure to make correct use of citations throughout, and have clarified vector notation as well as made sure to define notation and exact values where relevant, including in all the places you identified.

Validity of the experimental setting

As we understand it, you're saying here that the learning during overtraining may not take place for odors with larger overlap with the target, for instance in our Figure 6. This is consistent with our framing of the phenomenon as margin maximization between target and nontarget, which have zero overlap and are thus far apart. If the probes are closer in representational space to the target (very high overlap) than the nontarget, they will never be separated by a hyperplane equidistant between the target and nontarget cluster centers. This is a feature, not a bug, of our framework. Our goal here to model the neural data, where the mice are shown overlap-1 probe odors as the held-out test examples, so the increase in accuracy on these overlap-1 probe trials with overtraining is the object we seek to explain, and that our margin maximization model (as Figure 6 shows) indeed does do this.

Loss interpretation

We agree with your statement that accuracy and loss are distinct measures of performance and the former is a discontinuous metric. We also agree the training loss must be decreasing (even if very slowly or imperceptibly) even after accuracy saturates if the weights continue to move/test loss continues to fall, this is tautologically true by the definition of gradient descent. Your point here is exactly what we're trying to communicate: much work in the grokking literature in deep learning outlines how small changes in train loss at late-time (after training accuracy saturates) can lead to large changes in test accuracy (see [Kumar et al, 2023], [Lyu et al, 2023] for examples). In the neural data reanalysis, we show how despite no outward changes in mouse behavior, there is continued decrease on training loss during overtraining. This is one sense in which "mice grok" -- because we see it leads to stronger downstream accuracy.

Miscellaneous and Questions

We've addressed all of these minor aesthetic issues. Further, we've reworked the overtraining reversal section entirely to make it clearer and more detailed. We've focused on making the connection between overtraining, feature learning, and adaptation to reversal clear, both theoretically in a linear network whose dynamics we analyze but also empirically in terms of interpretable feature learning in the model.

We've also made further additions to the manuscript that we believe strengthen it. In addition to the reworked overtraining reversal section, the new manuscript we've uploaded (changes in blue):

• Demonstrates the nature of the objective is the key to this late-time learning, both through a loss ablation that shows a causal link between margin-maximization and grokking (Appendix A.6) but also by constructing a new, more biologically realistic model of piriform cortex (Appendix A.7), and showing the phenomenology persists there if the same loss function is used.
• Contains new statistical comparisons and ablations between the piriform separation during overtraining and random firing rate baselines in Appendix A.5.

Overall, we thank the reviewer for their feedback, and would appreciate if they would reconsider their assessment in light of our response and new additions to the manuscript. It seems the fundamental message of the paper was not appreciated, and much of the criticism was about minor aesthetic and prose details rather than about the core premise: proposing and modeling a new neural phenomenon using the rich language of machine learning. Thank you.

--

References.

Kumar, Tanishq, et al. "Grokking as the transition from lazy to rich training dynamics." arXiv preprint arXiv:2310.06110 (2023).

Lyu, Kaifeng, et al. "Dichotomy of early and late phase implicit biases can provably induce grokking." arXiv preprint arXiv:2311.18817 (2023).

Validity of the Experimental Setting

We think this issue might be a miscommunication/misconception. If the concerns stem from the impression that we are passing in raw n-hot vectors into the MLP, we should clarify this is not what we are doing. In fact, we already are explicitly doing something similar to what you suggest: passing in embeddings that are non-zero and continuous, into the network. The n-hot vectors are passed through a random projection before being fed into the network. We state this explicitly in the paragraph:

We take our earlier definition of odors as n-hot vectors of odorants (chemicals) that we pass through a random projection, inspired by the functional role of the glomeruli (Blazing & Franks, 2020), with n = 10, k = 100. We train a one-hidden layer multilayer perceptron (MLP) on these embeddings...

Loss Interpretation

You are right that there is obviously no notion of literal "training loss" for learning taking place in mouse piriform cortex, since the explicit form of the objective optimized by cortex is not known. Rather, what we are referring to in that quote is a) the continued separation of class representations of training odors that we present in Figure 2 and Figure 3 as well b) the continued increase in margin of the associated representations, that we present in Figure 4. As we can see in the PCA diagrams in Figure 5, the decreased training loss is a direct consequence of the separation of class representations in the synthetic setting. It is this separation (and associated increase in margin) that we are referring to as close proxies for "training loss" during mouse overtraining.

Thank you for engaging with our response, and please let us know what we can do to make it easier for you to consider reassessing your score.

The random projection is a multiplication of the sparse input vectors representing odors by a matrix whose entries are sampled from a unit Gaussian, then kept fixed throughout. So the input vectors to the MLP are of the form $y = Wx$ where $x$ is an $n$-hot vector representing a fixed odor and $W$ a fixed random matrix. The key point is these vectors $y$ will have nonzero overlap almost surely, even if the expected overlap is zero. Importantly, this is the point of the model -- in the experimental setting of [Berners-Lee et al, 2023], zero overlap odors are passed into the mouse olfactory cortex, whose early layers involve random projection.

We also feel the reviewer is focusing on the wrong details. The synthetic model, as reviewer CDUr notes, captures key computational principles at an abstract level of the main findings of the paper which are Figures 2, 3, 4. In response to reviewers, during the rebuttal period we added a new model that explicitly studies "realistic computational neuroscience," taking olfactory cortex biology much more seriously. This model, in Appendix A.7, is a literal interpretation of piriform anatomy, including modeling differing cell types, recurrent feedback, sparsity constraints, and more. We find given the same objective, the late-time learning in test loss is preserved.

Thank you for the suggestion to reword "training loss" and address the citet note. We will make both changes for the camera-ready version, as we cannot currently edit the manuscript. We hope during discussion the reviewer engages with other reviewers to discuss and appreciate the central finding of the paper: Section 3 demonstrates mice are continuing to learn even as their behavior remains unchanged, leading to abrupt generalization downstream, and models of grokking in deep learning make extremely non-obvious predictions (margin-maxization) about piriform data that turn out to be true! We think this is exceptionally interesting, and do appreciate the reviewer's willingness to engage in discussion even if we feel they have not appreciated the central contents of the paper.

Thank you for the detailed and positive review. We're rephrased the description of the synthetic model in the main text, adopting your excellent description, added the relevant citations to the literature you brought up, and reworked the explanation of margin maximization, in direct response to your suggestions.

We've also made a number of substantive additions to the manuscript we believe make it stronger and clearer. In particular, the updated manuscript we have uploaded contains changes in blue, and:

• Demonstrates the nature of the objective is the key to this late-time learning, both through our Hinge loss ablation but also by constructing a new, more biologically realistic model of piriform cortex (Appendix A.7), and showing the phenomenology persists there if the same loss function is used.
• Includes discussion for what biological analogs of such an objective might be in cortex, relating to existing work on how uncertainty is stored and used to guide behavior in rodent cortex.
• Contains new statistical comparisons and ablations between the piriform separation during overtraining and random firing rate baselines in Appendix A.5.
• Contains an entirely reworked overtraining reversal section to include a worked out mathematical example of how features can be re-used under label inversion, making it both clearer and more detailed.

We'd love to hear anything we can do to further improve your assessment of our work, and thank you for the detailed and thoughtful review!

Thanks for the review, and close read! As we understand it, there are three explicit points to be addressed. It's most natural to address your questions in reverse order. Our updated manuscript contains changes in blue throughout.

3. Demonstrate a causal link between margin maximization and the emergence of grokking in the MLP model?

Certainly. A natural way to ablate margin is to change the objective from one known to drive asymptotic margin-maximization to another that drives margin-maximization up to a certain point, after which the loss vanishes. This is a "hard margin" objective, and Hinge loss is an example. In the new Appendix A.6, we repeat our experiment with Hinge loss and margin $C=1$ and find when we ablate late-time margin maximization, test loss plateaus when train loss does, and stops improving at late time so that the grokking effect vanishes. We can also see in the accompanying PCA plots that the separation between target and probes is worse in that setting as a result. This strongly suggests margin maximization is driving late-time learning in our setting.

2. How results from MLP could generalize to more complex DL models for classification?

We want to emphasize our claim is not a universality claim about margin maximization driving grokking in all possible deep learning settings. Our focus is exclusively on a) piriform cortex, and b) our synthetic model thereof. We aim to focus on answering the question of how much such a phenomenology describes the neural data we revisit. While it's certainly plausible one might find a setting where, for instance, a vision model groks on CIFAR in such a manner, we think in this case looking for/including such a setting would actively distract from our topic of study, which is a deep-learning motivated analysis of overtraining dynamics in piriform cortex.

1. How results from the MLP model would generalize to more biologically constrained models of the PPC (or other olfactory processing circuits)?

We can check this directly. This offers us another chance to try falsify our claim that the objective, not architecture, is key here when it comes to margin maximization. In the new Appendix A.7, we consider an MLP architecture that takes posterior piriform biology much more seriously, modeling the different cell types as different layers with different initializations, their sparse vs dense/excitatory vs inhibitory effects with dropout and differing signs on the forward pass, as well as recurrent feedback effects with a leaky ReLU. We find this model -- when trained with cross-entropy -- exhibits the same phenomenology in the same way, supporting the idea that the task and objective, not the details of the architecture, are what are key here. We elaborate on what the corresponding objective might be in cortex by discussing related work from rodent neuroscience in the Discussion section.

We emphasize, both here and in our updated manuscript, that truly the only way to check whether the margin-maximization hypothesis is an accurate description of late-time learning in cortex is to run new and targeted experiments in sensory cortex itself, rather than construct even more machine-learning-esque models. To this end, we lay out -- in detail -- an experimental proposal and our falsifiable hypotheses, and a plan for how this can be done in the future, in Appendix Section A.2, to make it easier for experimentalists to test our hypotheses.

Other substantive additions to the manuscript:

• A reworked overtraining reversal section to include a fleshed out mathematical example of how features can be re-used under label inversion, making it both clearer and more detailed, with both theory and experiments.
• Statistical comparison of representational separation that takes place in cortex during overtraining to random/control baselines in the new Appendix A.5.

We think overall both the additions specifically in response to your questions, and otherwise, significantly strengthen the paper. We'd appreciate it if you'd consider raising your score, or otherwise suggest additional things we can do for you to do so.

Thank you for your thoughtful review; we'd love feedback on our substantial additions and many new experiments done directly in response to your comments, and would appreciate if you consider raising your score as the rebuttal period comes to a close. Thank you!

Thanks for the review! Your aesthetic notes have all been addressed in the updated manuscript we've uploaded -- changes in blue text throughout. Note the target dot being obscured was not an accident, part of the point being made is that target-probes are not easily separable at the end of training, but become separable during overtraining. Overtraining starts at the beginning of Fig 1, and we've put that in the plot title so it's front and center, as well as included column titles (mice names) + cluster centers in Figs 2 and 3, respectively.

Your question about a non-example is a good one. First, to directly address the question we've added a specific non-example: ablating cross-entropy to hinge loss in the MLP experiment. Hinge loss penalizes hard margin, so that after all points are classified correctly with some margin $C$, loss vanishes, in contrast to loss-entropy, where it doesn't vanish. We see the phenomenon vanishes concomitantly, strongly suggesting margin maximization plays a causal role in driving grokking in our model. These new experiments and ablations can be found in Appendix A.6.

In addition to addressing your aesthetic comments and demonstrating a causal link to address your question, our updated manuscript includes much more content. It now

• Demonstrates the nature of the objective is the key to this late-time learning, both through our Hinge loss ablation but also by constructing a new, more biologically realistic model of piriform cortex (Appendix A.7), and showing the phenomenology persists there if the same loss function is used.
• Includes discussion for what biological analogs of such an objective might be in cortex, relating to existing work on how uncertainty is stored and used to guide behavior in rodent cortex.
• Contains new statistical comparisons and ablations between the piriform separation during overtraining and random firing rate baselines in Appendix A.5.
• Contains a reworked overtraining reversal section to include a worked out mathematical example of how features can be re-used under label inversion, making it both clearer and more detailed.

We think including including substantially more content while simultaneouly polishing the figures and prose strengthens the submission significantly. We'd appreciate it if you considered raising your score as a result, or pointed us to further desiderata for the manuscript. Thanks!