Biologically Inspired Mechanisms for Facilitating Grokking in Multilayer Perceptrons

Motivation

• Grokking is delayed generalization in which near-perfect training performance precedes a later transition toward systematic, rule-like solutions and reorganized internal subnetworks.This phenomenon has been studied mainly on small algorithmic tasks (Power et al., 2022), with internal changes linked to generalizable subnetworks (Merrill et al., 2023).Source: §1 · p. 1
• Conventional MLP training minimizes prediction error without directly encouraging modularity, specialization, sparse circuits, or robust reusable internal patterns.The fractured entangled representation hypothesis argues that objective-driven training can distribute useful regularities across redundant, interfering components (Kumar et al., 2025).Source: §1 · pp. 1–2

Broader relevance

• Because large language models also reorganize internal representations during extended training, regulating internal organization may accelerate systematic behavior and reduce training cost.The paper uses grokking as a setting for studying how architectural biases shape representations during learning.Source: §1 · p. 2

Benchmarks

• The study evaluates sparse parity and noisy XOR-cluster classification, which expose delayed generalization through compact hidden dependencies and nonlinear geometry-aligned features, respectively.Sparse parity uses binary vectors with labels determined by a small hidden coordinate subset (Merrill et al., 2023), whereas noisy XOR uses four Gaussian clusters with some flipped labels (Xu et al., 2023).Source: §1 · p. 2

Contribution

• BioNN incorporates biologically inspired hidden-layer mechanisms that regulate activity routing, neuronal competition, and effective connectivity during learning.The mechanisms include input gating, structural plasticity, gain modulation, threshold modulation, homeostasis, lateral inhibition, and activation decorrelation.Source: §1 · p. 2

2.1 Grokking

• Grokking accounts emphasize optimization and regularization, delayed feature learning, internal reorganization, and geometric or physical perspectives, but no unified explanation has yet emerged.Weight decay can slowly move models from high-norm training-fitting solutions toward lower-norm ones, while feature learning can produce grokking without weight decay; existing accounts often operate at different analytical levels.Source: §2.1 · pp. 3, 4, 6
• Studies link grokking to sparse circuit emergence, task-specific feature specialization, relative flatness, higher-entropy states, and interventions that shorten delayed generalization.Merrill et al. (2023) describe competition between dense memorizing and sparse generalizing circuits, while Han et al. (2025) find relative flatness changes sharply at generalization; Lee et al. (2024) accelerate generalization by amplifying slow gradient components.Source: §2.1 · p. 5

2.2 Brain-Inspired Neural Models

• Brain-inspired neural models impose internal organization through neuron- and synapse-level mechanisms, local learning, predictive coding, richer dynamics, structured connectivity, and specialized representations.These approaches include spiking neural networks, predictive coding, dendritic nonlinearities, oscillatory computation, synchronization, and reservoir-based dynamics, despite differing in biological inspiration and implementation.Source: §2.2 · pp. 6, 8
• Existing models demonstrate biological inspiration through online local plasticity, learned sparse connectivity, incremental predictive coding, continual learning, nonlinear dendritic integration, and wave-based computation.Examples include e-prop’s local eligibility traces and learning signals, learned recurrent connection probabilities, and incremental predictive coding that updates parameters throughout inference.Source: §2.2 · pp. 6–7

3.1 Notation List

• The notation list identifies the most relevant parameters of the BioNN model.The paper notes that the model contains a substantial number of parameters and therefore summarizes the most relevant ones in Table 1.Source: §3.1 · p. 8, Table 1 · p. 8
• The parameter notation is organized as a reference for the BioNN model description.Source: §3.1 · p. 8, Table 1 · p. 8
• The supplied section contains parameter-list material rather than reported experimental findings.Source: §3.1 · p. 8, Table 1 · pp. 8–11
• Table 1 presents the initial portion of the model-parameter list.Source: Table 1 · p. 8
• The notation coverage spans multiple continued tables rather than a single table segment.Source: Table 1 · pp. 8–11
• Table 1 continues across additional table sections to accommodate the full parameter notation.Source: Table 1 · pp. 9–11

• The model retains the MLP hidden-layer role of transforming inputs into representations, but regulates signal entry, responses, connectivity, competition, and activity stability.These mechanisms modify rather than replace the core computation.Source: §3.2 · p. 11
• The input-derived context representation controls input gating, gain modulation, and fast threshold modulation, while other mechanisms operate independently of context.Context-dependent regulation is sample-specific and can amplify or suppress input coordinates.Source: §3.2 · p. 12
• The regulated forward pathway applies masked sparse connectivity, gain and threshold modulation, nonlinear activity, and lateral inhibition before producing the final hidden activity.Lateral inhibition suppresses diffuse co-activation among hidden neurons.Source: §3.2 · pp. 12–13
• Training-time regulation comprises homeostatic threshold updates, structural plasticity of active connections, and activation decorrelation, whereas lateral inhibition operates during forward computation.Homeostasis adjusts long-term thresholds, structural plasticity changes effective connectivity, and decorrelation discourages repeated encoding of the same information.Source: §3.2 · pp. 12–13

3.3 Input Gating

• The input gate reweights input features according to context before the main hidden-layer computation, changing the effective input representation received by hidden neurons.It increases some feature contributions and reduces others, allowing the same hidden neuron to respond differently to context-adjusted inputs.Source: §3.3 · pp. 13–14
• The gate provides a compact neural-network abstraction of dendritic pathway-specific modulation and selective integration.This design reflects biological dendrites’ structured processing, including nonlinear integration and local adjustment of inputs from different sources.Source: §3.3 · p. 14
• Formally, the gate rescales each input coordinate before the affine transformation, with g containing one context-dependent scaling factor per input dimension.Values of g greater than one amplify inputs, whereas values smaller than one attenuate them, producing a reweighted input xg.Source: §3.3 · p. 14
• The gate is initialized neutrally so the input is initially unchanged, then learns context-dependent reweighting through backpropagation.The trainable gate parameters are G and bg, while C and bc define the shared context representation; the gate head starts at zero, yielding the neutral vector 1.Source: §3.3 · p. 14

3.4 Structural Plasticity

• Structural plasticity reflects experience-dependent reconfiguration in which useful connections stabilize while weak or interfering pathways are removed or reorganized.This biological motivation links connectivity changes to adaptation, memory, and specialized computation over multiple timescales.Source: §3.4 · pp. 14–15
• A binary structural mask M selects effective connections, while the dense trainable matrix W remains stored and only masked entries participate in the forward pass.Mij = 1 activates connection i,j; Mij = 0 removes it from the effective circuit.Source: §3.4 · p. 15
• The mask controls active density ρm and can be periodically refreshed by retaining the top-k entries of |Wij|, with k = max(1, ⌊ρmnd⌋).This keeps large-magnitude connections in the effective circuit while removing weak ones.Source: §3.4 · p. 15
• When ρm < 1, periodic updates select among sparse effective circuits; when ρm = 1, all connections remain active, isolating sparsification from the plasticity mechanism itself.W is optimized by gradient descent, whereas M is updated explicitly outside gradient descent according to the experiment’s schedule.Source: §3.4 · pp. 15–16

Biological motivation

• Gain modulation changes response amplitude or sensitivity while preserving a neuron's basic selectivity, making it a major computational principle in the central nervous system.Biological gain is influenced by attention, arousal, motor state, local circuit conditions, and neuromodulators such as acetylcholine and noradrenaline.Source: §3.5 · p. 16

Formal mechanism

• The model applies gain modulation to each hidden neuron's preactivation through a per-neuron vector γ, with α_g controlling modulation strength.The vector γ determines each neuron's current excitability, allowing the same representational components to contribute with different strengths across contexts.Source: §3.5 · p. 16

Implementation

• Large gain increases a neuron's response to the same input, whereas small gain decreases it; parameters A and b_a are learned by backpropagation.The gain head is initialized to zero, producing the neutral initial gain vector 1.Source: §3.5 · p. 16

3.6 Threshold Modulation

• Biological threshold modulation enables neurons to become easier or harder to activate based on recent state, input dynamics, and network activity.Rapid input accumulation can lower firing criteria, while adaptive changes preserve sensitivity to rapid fluctuations and reduce slow-input effects.Source: §3.6 · pp. 16–17
• The model separates fast sample-dependent threshold shifts from slower activity-dependent homeostatic threshold changes.The fast component varies through the context vector, whereas the homeostatic component changes more slowly through activity-dependent updates.Source: §3.6 · p. 17
• Learned threshold-modulation parameters are initialized so the fast threshold shift begins at zero, giving current input direct control over hidden-neuron activation.The parameters T and b_t are learned by backpropagation, and the threshold-modulation head is initialized to produce Δθ = 0.Source: §3.6 · p. 17

3.7 Homeostasis

• Homeostatic plasticity stabilizes neural activity by lowering responsiveness when activity is high and increasing it when activity is low.Synaptic scaling adjusts overall excitatory synaptic strength while preserving much of the learned relative pattern.Source: §3.7 · p. 17
• The model assigns each hidden neuron a slow threshold state that rises when activation is too frequent and falls when activation is too rare.This threshold state regulates response frequency while gradient-based learning organizes neuronal selectivity, matching a multi-timescale learning framework.Source: §3.7 · p. 18
• The update uses minibatch activation frequency, an activity cutoff, a target activity rate, an update rate, and a bound on threshold magnitude.The threshold vector is a state variable rather than a parameter learned by gradient descent.Source: §3.7 · p. 18
• Explicit updates from final hidden activity counteract collapse into a few persistently active neurons and many persistently silent neurons.The update is performed without gradient descent.Source: §3.7 · p. 18

Biological motivation

• Lateral inhibition sharpens representations by suppressing competing activity, supporting selective responses, winner-take-all behavior, and reduced diffuse activation.Its contribution to stimulus selectivity may overlap with effects from feedforward structure, spike threshold, and response nonlinearities.Source: §3.8 · pp. 18–19

Model implementation

• The mechanism applies lateral inhibition after gain and thresholding, using each neuron’s mean activity across the other hidden neurons as global inhibitory pressure.This competition makes diffuse activation harder to maintain and selective patterns easier to stabilize.Source: §3.8 · p. 19

Implementation details

• The model uses no separate learned inhibition matrix: β is fixed by the experimental configuration, while gradients propagate through the inhibition computation to trainable layer parameters.β controls the strength of lateral inhibition.Source: §3.8 · p. 19

3.9 Activation Decorrelation

• Biological motivation comes from evidence that retinal processing, adaptive dynamics, and inhibition reduce redundancy or shared correlations in neural responses.These mechanisms can attenuate predictable input, reduce temporal correlations, and promote sparse, decorrelated representations.Source: §3.9 · pp. 19–20
• Activation decorrelation encourages hidden neurons to specialize rather than repeatedly encode the same variation across instances.The regularizer penalizes similar firing patterns across a minibatch, pushing the model toward less redundant hidden codes.Source: §3.9 · p. 20
• The model computes decorrelation from centered, column-normalized final hidden activations and removes self-correlations before applying the loss.The decorrelation strength is controlled by λd, and the penalty is added to the task loss for gradient-based training.Source: §3.9 · p. 20

4.1 The Grokking Problems Under Study

• The study uses sparse parity and noisy XOR classification to distinguish memorization from generalization in problems with simple rules embedded in irrelevant variation (Merrill et al., 2023; Xu et al., 2023).Both tasks permit finite-sample fitting through accidental correlations, whereas generalization requires recovering the underlying compact structure.Source: §4.1 · pp. 20–21

4.1.1 Sparse Parity

• Sparse parity uses 40-dimensional sign vectors whose label is the product of the first three coordinates, leaving 37 random distractor coordinates.Its non-additive three-way interaction makes it suitable for testing whether networks learn a compact rule rather than dense sample-specific memorization.Source: §4.1.1 · p. 21

4.1.2 Noisy XOR Clusters

• Noisy XOR uses four latent clusters in R40000, with class determined by the signal axis, isotropic noise, and 5% independently flipped training labels while test labels remain clean.The task tests whether representations organize around latent cluster geometry despite nonlinear class structure, high dimensionality, and corrupted supervision.Source: §4.1.2 · pp. 21–22

4.2 Network Structure Indicators

• The trained network’s hidden activations are evaluated using lifetime sparsity, sample Hoyer sparsity, neuron Hoyer sparsity, dead-neuron fraction, and singular-value entropy.Activity is defined by |Hiu| > ε, with ε = 10−4 for sparse parity and ε = 10−7 for noisy XOR; groups contain two target classes or four latent clusters, respectively.Source: §4.2 · pp. 22–24

4.3 Experimental Protocol

• Both tasks use a one-hidden-layer BioNN with 1000 neurons, matched ablations and random seeds, while SGD protocols differ by task and each configuration has 20 independent runs.Sparse parity trains for 300 epochs with hinge loss; noisy XOR trains for 5000 steps with logistic loss, and additional Adam tests produced generally worse or absent generalization within the evaluation period.Source: §4.3 · p. 27

5.1 Ablation Results

• Homeostasis and structural plasticity dominate the ablations, although their relative contributions differ between sparse parity and noisy XOR.Successful generalization is consistently associated with regulating neuron use or effective connectivity, but neither mechanism is uniformly sufficient across both tasks.Source: §5.1 · p. 35

5.2 Discussion: Sparse Parity

• In sparse parity, homeostasis plus structural plasticity nearly reproduces the full model, reaching Train 1 = 19 and Test 1 = 22 versus the all-off baseline’s Test 1 = 137.Homeostasis alone reaches Test 1 = 69, while the strongest configurations have about 7.2–7.4 active neurons.Source: §5.2 · pp. 36–38

5.3 Discussion: Noisy XOR

• In noisy XOR, homeostasis is most reliable under dense connectivity, but structural sparsification at ρm = 0.35 enables generalization without homeostasis while usually delaying perfect convergence when homeostasis is active.At ρm = 1, homeostasis alone reaches Test 1 = 2571; at ρm = 0.35, all twelve repeated no-homeostasis configurations reach Test 0.9 and seven reach perfect test accuracy.Source: §5.3 · pp. 39–40

5.4 Mechanism-by-Mechanism Interpretation

• Mechanism-by-mechanism, gain modulation complements homeostasis in XOR, whereas input gating, threshold modulation, lateral inhibition, and activation decorrelation are not individually required for generalization.Lateral inhibition’s effect may differ under richer recurrent inhibitory dynamics than the simplified implementation used here.Source: §5.4 · pp. 42–43

5.5 Implications for Generalization and Large Language Models

• The findings suggest that homeostasis and adaptive structural sparsification could regulate component use and effective architecture in large language models, with benefits determined by task and surrounding mechanisms.Possible targets include neurons, channels, attention heads, experts, adapters, and blocks; structural adaptation should be treated as context-dependent capacity allocation rather than uniformly beneficial sparsity.Source: §5.5 · p. 43
• Lateral inhibition, activation decorrelation, gain modulation, and threshold modulation appear secondary to homeostasis and structural plasticity in the present ablations.Decorrelation is not required for differentiated representations and is better viewed as a complementary objective, while lateral inhibition remains relevant where specialization and competition are central.Source: §5.5 · p. 44
• Homeostasis, structural plasticity, input gating, and lateral inhibition address distinct aspects of allocating computational capacity as useful internal structure develops.The experiments suggest capacity allocation may matter for generalization as much as the amount of available capacity.Source: §5.5 · p. 44

5.6 Computational Demands

• Per-step overhead does not capture the full computational trade-off because mechanisms adding computation may still be favorable if they reduce optimization steps to high test accuracy.The ablations show large differences in the time required to reach high test accuracy.Source: §5.6 · p. 44
• Homeostasis and simplified lateral inhibition have O(Bn) overhead, while structural plasticity updates masks periodically and can reduce inference cost only with efficiently exploitable sparsity.Homeostasis requires O(n) memory; masked dense structural representations retain most dense-computation cost.Source: §5.6 · pp. 44–45
• Activation decorrelation requires O(Bn^2) computation and O(n^2) memory, making it difficult to scale directly, whereas context-dependent modulation adds projections whose cost depends on modulation dimension.In the present model, the context dimension is s = 32; modulation projections remain active at inference, while decorrelation is primarily a training-time operation.Source: §5.6 · p. 45

• Homeostasis and structural plasticity substantially accelerate memorization-to-generalization transitions, enabling generalization in configurations that otherwise remain in memorization.Homeostasis most consistently regulates long-term hidden-neuron use, while structural plasticity changes which connections participate in effective computation.Source: §6 · p. 46
• Grokking reflects internal reorganization alongside optimization, involving changes in neuron recruitment, computation distribution, and functionally important connections.In sparse parity, the mechanisms produce a compact sparse circuit; in noisy XOR, successful generalization involves greater hidden-representation sparsity, specialization, and latent-structure separation.Source: §6 · p. 46
• The BioNN implementation provides a reusable testbed with independent switches for controlled comparisons of seven biologically inspired mechanisms and their combinations.Source: §6 · p. 46
• For LLMs, homeostatic control and structured plasticity may promote earlier stable representations, adaptive sparse computation, and fewer optimization steps for generalization.This possibility remains a question beyond the present experiments; homeostasis is attractive because its computational overhead is small.Source: §6 · p. 46