PaperlayerView exact revision on arXiv

Source-linked AI summary

Biologically Inspired Mechanisms for Facilitating Grokking in Multilayer Perceptrons

Florin Leon

arXiv:2608.28184v1cs.LG

TL;DR

Grokking involves delayed generalization and internal reorganization, motivating the question of whether biologically inspired regulation can promote this transition. The paper augments an MLP with seven such mechanisms and evaluates them on sparse parity and noisy XOR. Homeostasis and structural plasticity show the strongest benefits, supporting explicit regulation of neuron utilization and effective connectivity within the studied tasks.

  • Problem

    The paper asks whether biologically inspired mechanisms can promote grokking by regulating hidden-layer activity, responses, and effective connectivity.

  • Method

    The study augments a multilayer perceptron with seven biologically inspired mechanisms and evaluates them through ablations on sparse parity and noisy XOR.

  • Results

    Homeostasis provides the strongest and most consistent benefit, while structural plasticity is the second major mechanism across the two grokking benchmarks.

  • Takeaways & Limitations

    The experiments support explicit regulation of neuron utilization and effective connectivity as a way to improve the emergence of generalizable internal computation.

  • Takeaways & Limitations

    The conclusions are qualified by the simplified lateral-inhibition implementation, which lacks explicit neuron-to-neuron inhibitory connections and recurrent competitive dynamics.

Abstract

from arXiv · show

Grokking is a delayed transition from memorization to generalization that is often accompanied by substantial reorganization of internal representations. This paper studies whether biologically inspired mechanisms, many of which are not commonly incorporated into artificial neural networks, can actively promote this transition by regulating hidden-layer computation at the levels of neuronal activity, response, and effective connectivity. We augment a multilayer perceptron with input gating, structural plasticity, gain modulation, threshold modulation, homeostasis, lateral inhibition, and activation decorrelation, and evaluate these mechanisms through systematic ablations on two established grokking benchmarks: sparse parity and noisy XOR classification. The results show that the mechanisms contribute unequally to generalization. Homeostasis provides the strongest and most consistent benefit, while structural sparsification emerges as the second major mechanism. The remaining biologically inspired mechanisms have smaller or less consistent effects in the present experiments. For both problems, the results support the common principle that explicit regulation of neuron utilization and effective connectivity can improve the emergence of generalizable internal computation. These findings motivate broader investigation of biologically inspired activity regulation and adaptive sparsification, including in large language models, where they may accelerate the development of generalizable representations and reduce the optimization time required for robust generalization.

1 Introduction

The introduction frames grokking as delayed generalization linked to internal reorganization, then motivates testing biologically inspired regulation in MLP hidden layers across sparse parity and noisy XOR.

  • Grokking is a delayed shift from near-perfect training performance with poor test performance toward later generalization and more systematic computation.
  • Conventional MLP training minimizes prediction error without directly enforcing modularity, specialization, sparse circuits, or reusable internal patterns.
  • The paper studies sparse parity and noisy XOR, two benchmarks exhibiting distinct delayed-generalization transitions involving compact circuits or changes beyond sample-level memorization.
  • BioNN incorporates biologically inspired mechanisms that regulate hidden-layer activity, neuronal competition, routing, and effective connectivity during learning.
  • The study asks whether these mechanisms improve the speed, stability, or form of the transition from memorization to generalization.

2 Related Work

Related work explains grokking through optimization, representation learning, geometry, physical analogies, and timescale separation, while brain-inspired models motivate imposing structure on learning and internal organization.

  • Existing grokking accounts emphasize optimization and regularization, delayed feature learning, internal reorganization, or geometric and physical descriptions.
  • Delayed feature-learning accounts connect grokking to representation change and alignment between initial neural tangent features and the target function.
  • Noisy XOR demonstrates that memorization of corrupted labels and near-optimal clean-test performance can emerge at different times.
  • Sparse parity studies describe grokking as replacement of a dense memorizing subnetwork by a smaller task-relevant circuit.
  • Neural collapse can precede generalization, whereas relative flatness changes sharply at generalization and appears more closely related to the transition.
  • Several theoretical and optimization studies model grokking as slow movement among training-fitting solutions or separation between fast overfitting and slower generalization dynamics.
  • No unified explanation has yet connected the different analytical levels used across grokking accounts.
  • Brain-inspired models motivate adding structure through signaling changes, local learning, competition, synchronization, and richer neural interactions.

3 BioNN Model Description

The supplied passages identify Table 1 as a list of model parameters presented across multiple continued tables, but do not describe individual parameters.

  • Table 1 lists model parameters for the BioNN implementation.
  • The supplied passages provide captions and implementation context but no parameter names, values, or functional descriptions.
  • The parameter list continues across additional table segments.

3.2 Overview of the Bio-Inspired Mechanisms

The BioNN retains the basic MLP hidden-layer computation but regulates input routing, neuronal responses, connectivity, competition, and activity stability through biologically inspired mechanisms. A compact input-derived context drives selected fast mechanisms, while other mechanisms adapt during training.

  • Overview of the Bio-Inspired Mechanisms: BioNN preserves the MLP hidden layer’s role while regulating how signals enter, how neurons respond, which connections remain active, and how activity is stabilized.The mechanisms modify, rather than replace, the core hidden-layer computation.
  • Overview of the Bio-Inspired Mechanisms: A compact context representation derived from the input controls input gating, gain modulation, and fast threshold modulation.Other mechanisms operate independently of this context.
  • Overview of the Bio-Inspired Mechanisms: The regulated layer first reweights inputs and applies sparse connections, then regulates hidden responses with gain and thresholds before computing final activity.The final hidden activity is passed to the readout layer.
  • Overview of the Bio-Inspired Mechanisms: Homeostasis adjusts long-term activation thresholds, structural plasticity changes effective connectivity, and decorrelation discourages repeated encoding of the same information during training.Lateral inhibition acts on primary activity before final hidden activity is computed.

3.3 Input Gating

Input gating provides context-dependent reweighting of input features before the main hidden-layer computation. It can amplify or attenuate coordinates so the same hidden neuron receives an adjusted version of the input.

  • Input Gating: The input gate applies context-dependent reweighting to input features before the main hidden-layer computation.This design abstracts pathway-specific modulation of incoming signals.
  • Input Gating: The gate increases some feature contributions and reduces others, changing the effective input representation received by hidden neurons.The same hidden neuron can therefore produce different responses under different contexts.
  • Input Gating: Input gating rescales each input coordinate, with values above one amplifying dimensions and values below one attenuating them.The resulting reweighted input depends on the context.
  • Input Gating: The gate is initialized to the neutral vector 1, leaving the input unchanged initially before learning context-dependent reweighting.Its trainable parameters are learned through backpropagation.

3.4 Structural Plasticity

Structural plasticity models learning as adaptation of effective connectivity, allowing the network to prune, preserve, and reorganize the connections used in forward computation. A binary mask selects active connections and can be periodically refreshed using weight magnitude.

  • Structural Plasticity: Structural plasticity represents continuous reconfiguration of connectivity across development, learning, and adult cortical adaptation.This biological process motivates changing effective connectivity during model training.
  • Structural Plasticity: The model uses a binary structural mask to determine which weight-matrix connections participate in effective computation.The dense trainable matrix is retained, but the forward pass uses only mask-selected entries.
  • Structural Plasticity: The active connection density is controlled by ρm, and the mask can be refreshed by retaining the top-k entries of |Wij|.Here k = max (1, floor(ρmnd)).
  • Structural Plasticity: When ρm < 1, only a fraction of connections participates in forward computation and periodic updates can change the active circuit.When ρm = 1, all connections remain active, providing a dense limiting case.
  • Structural Plasticity: Weights are optimized by gradient descent, whereas the structural mask is updated separately by an explicit periodic rule.Thus learning operates over both parameter values and circuit structure.

3.5 Gain Modulation

Gain modulation regulates the strength of hidden-neuron responses without changing their basic selectivity. In BioNN, context-dependent gains alter neuronal excitability for each instance while preserving the underlying representational components.

  • Gain Modulation: Gain modulation changes response amplitude or sensitivity while preserving a neuron’s basic selectivity for preferred input patterns.Biological gain can vary with attention, arousal, motor state, or neuromodulatory input.
  • Gain Modulation: The mechanism abstracts state-dependent excitability, allowing the same basic circuit organization to process signals with different amplification levels.Gain changes response strength within the current context.
  • Gain Modulation: Gain modulation changes how strongly each hidden neuron responds to its preactivation.The gain vector has one entry per hidden neuron and is controlled by αg.
  • Gain Modulation: A large gain strengthens a neuron’s response to the same input, whereas a small gain reduces responsiveness.The gain parameters are learned by backpropagation.
  • Gain Modulation: The gain head starts at zero, producing the neutral gain vector 1 at initialization.This lets the network initially preserve the unmodulated response before learning context-dependent amplification.
  • Gain Modulation: Gain modulation lets the network retain the same representational components while varying their contribution across contexts.The mechanism changes response strength rather than replacing the components.

3.6 Threshold Modulation

Threshold modulation makes hidden-neuron activation depend on both the current input context and a slower homeostatic state. The mechanism adjusts how readily neurons activate in response to changing evidence.

  • Biological motivation: Biological threshold modulation allows the same input pattern to produce different outputs as a neuron's internal activation threshold changes.Thresholds vary with recent cell state, incoming depolarization, background activity, and intrinsic excitability.
  • Biological motivation: Rapidly accumulating input can lower the effective firing criterion, while adaptive changes preserve sensitivity to rapid fluctuations despite slow input changes.
  • Model mechanism: The hidden neuron uses a total threshold combining fast context-dependent modulation with a slower activity-dependent homeostatic component.The fast shift changes across samples through context c, whereas the homeostatic threshold changes more slowly.
  • Model mechanism: The threshold-modulation head starts with no fast shift, then lets the current input directly influence how easily each hidden neuron activates.The threshold parameters are learned by backpropagation, and initialization sets ∆θ = 0.

3.7 Homeostasis

Homeostasis stabilizes hidden-layer activity by adjusting each neuron's threshold according to its recent activation frequency. This slow feedback process counters excessive activity and near-silence while learning continues.

  • Biological motivation: Homeostatic plasticity keeps neural activity within a stable range by reducing responsiveness when activity is high and increasing it when activity is low.The biological principle can operate through mechanisms such as synaptic scaling while preserving much of learned synaptic structure.
  • Computational principle: The computational division is that learning shapes neuronal selectivity, while slower regulation controls how often each neuron responds.
  • Model mechanism: The model assigns each hidden neuron a slowly changing threshold that rises when activation is too frequent and falls when activation is too rare.The update acts on threshold state rather than the weight vector defining selectivity.
  • Update rule: For each minibatch, the model measures the activation frequency of every hidden neuron before updating its homeostatic threshold.The activity measure uses the minibatch activations and an activity cutoff ε.
  • Update rule: The update rate, target activity rate, and threshold bound control the speed, goal, and magnitude of homeostatic adaptation.
  • Functional effect: Explicit threshold updates counter population collapse into a few persistently active neurons and many persistently silent neurons without using gradient descent.

3.8 Lateral Inhibition

Lateral inhibition introduces competition among hidden neurons so that strongly supported responses become more selective while weaker or redundant co-activations are reduced. In this model, inhibition is applied globally after gain and threshold processing.

  • Biological motivation: Lateral inhibition sharpens population responses by suppressing neighboring or competing activity and emphasizing the strongest response.The biological effect is proposed as one contributor to selectivity alongside feedforward structure, thresholds, and nonlinearities.
  • Computational principle: The competitive hidden layer favors a cleaner allocation in which the most strongly supported response dominates and similar alternatives are reduced.This also limits broad, redundant co-activation among similar neurons.
  • Model mechanism: The model applies lateral inhibition after gain and thresholds have produced a candidate hidden activity.
  • Model mechanism: Each hidden neuron receives inhibitory pressure based on the mean primary activity of all other hidden neurons, making inhibition global across the layer.
  • Model mechanism: The final hidden activation subtracts an inhibition term whose strength is controlled by β.
  • Functional effect: Competition makes diffuse activation harder to maintain and selective patterns easier to stabilize without adding a learned inhibition matrix.β is fixed by the experimental configuration, while gradients propagate through the inhibition computation.

3.9 Activation Decorrelation

Activation decorrelation targets redundancy in hidden representations by encouraging neurons to encode complementary variation rather than repeatedly representing the same signal. The model implements this principle as a loss on hidden activations.

  • Biological motivation: Efficient coding treats redundant population responses as a computational problem and favors neurons carrying complementary information.This motivation is especially relevant for sensory representations constrained by limited bandwidth.
  • Biological motivation: Retinal processing provides evidence for redundancy reduction because ganglion-cell spike trains are less correlated than corresponding naturalistic visual inputs.Retinal responses also attenuate predictable patterns and emphasize deviations from recent stimulus history.
  • Biological motivation: Adaptation can reduce temporal output correlations, while inhibition can reduce shared fluctuations across nearby neurons.
  • Computational principle: The proposed mechanism encourages hidden neurons to specialize instead of repeatedly encoding the same variation across instances.
  • Model mechanism: The model centers and normalizes minibatch hidden activations before constructing the hidden-neuron correlation matrix.The activation matrix has one row per example and one column per hidden neuron.
  • Model mechanism: Self-correlations are removed from the correlation matrix before the decorrelation penalty is computed.
  • Model mechanism: The decorrelation loss penalizes remaining hidden-neuron correlations, with λd controlling its strength.Because it is added to the task loss, gradients update the trainable model parameters.

4 Experimental Setup

The experiments use sparse parity and noisy XOR classification to study grokking under high-dimensional distractors, nonlinear structure, and label noise. Hidden-layer representations are evaluated with activity, sparsity, specialization, and neuron-utilization indicators, alongside optimizer comparisons.

  • Tasks: The study evaluates sparse parity and noisy XOR classification, two binary tasks whose compact rules are obscured by high-dimensional irrelevant variation.Sparse parity uses 40-dimensional sign vectors with three relevant coordinates; noisy XOR uses 40,000-dimensional inputs organized around four latent clusters with isotropic noise.
  • Tasks: Sparse parity tests whether a network shifts from dense, sample-specific memorization toward a computation based on the small subset of coordinates determining the label.The target is y = x1 · x2 · x3, while the remaining coordinates are random distractors; the relevant rule is a higher-order interaction rather than an additive one.
  • Tasks: Noisy XOR tests whether the model recovers cluster-level structure despite flipped training labels and many directions for fitting finite-sample fluctuations.Positive and negative classes each occupy two opposite clusters, and the experiments flip 5% of training labels while keeping test labels clean.
  • Network structure indicators: The hidden activation matrix is analyzed using lifetime sparsity, sample Hoyer sparsity, neuron Hoyer sparsity, dead-neuron fraction, and group-based structural indicators.These measures assess restricted neuron usage, concentration of activation magnitudes, effectively unused neurons, cluster separation, specialization, and support overlap.
  • Experimental protocol: The protocol also compares SGD with newer optimizers, with Adam yielding generally worse performance and different grokking dynamics in additional tests.The authors retain SGD because it was used in the original work or code associated with both grokking problems, while noting that optimizer effects require further study.

5 Results and Discussion

Homeostasis and structural plasticity account for most of the improvement in grokking across sparse parity and noisy XOR, but their effects depend on task and structural density. Together, they regulate neuron use and effective connectivity, promoting earlier or more reliable generalization.

  • Overall ablation pattern: Homeostasis and structural plasticity are the dominant mechanisms, although their relative roles differ between sparse parity and noisy XOR.Across both tasks, the strongest effects come from regulating neuron use and effective connectivity.
  • Sparse parity: Perfect test accuracy arrives at epoch 18 for the fully enabled parity model versus epoch 137 for the all-off baseline.The best parity configuration reaches perfect test accuracy at epoch 17 despite removing lateral inhibition and activation decorrelation.
  • Sparse parity: Top parity configurations consistently enable homeostasis and structural plasticity, while the bottom configurations consistently disable both.This pattern identifies the pair as nearly sufficient for the parity effect, with the two-mechanism configuration reaching Train 1 = 19 and Test 1 = 22.
  • Sparse parity: Removing structural plasticity delays parity Test 1 from 18 to 53, while structural plasticity alone reaches Test 1 = 43.Homeostasis alone reaches Test 1 = 69, making structural plasticity the larger individual contributor in this task.
  • Sparse parity: The strongest parity configurations compress computation into very small active subnetworks, averaging about 7.4 neurons in the best row and about 7.2 with both major mechanisms enabled.They also show higher activation and Hoyer sparsity, lower effective rank, and a dead-neuron fraction of about 0.713.
  • Noisy XOR: At ρm = 1, every XOR configuration reaching perfect test accuracy contains homeostasis, whereas at ρm = 0.35 structural sparsification enables generalization in some configurations without homeostasis.Without homeostasis at ρm = 0.35, all twelve repeated configurations reach Test 0.9 and seven reach perfect test accuracy; with homeostasis, sparsification reaches Test 0.9 earlier but Test 1 later in sixteen of seventeen matched cases.
  • Noisy XOR: XOR generalization lacks a single representational signature, because homeostasis and structural sparsification impose partly overlapping internal organizations whose interaction depends on context.Without homeostasis, reduced structural density produces sparser, less correlated representations alongside increased specialization and centroid separation.

6 Conclusions

The study finds that regulating neuron utilization and effective connectivity can substantially influence grokking, with homeostasis and structural plasticity producing the strongest effects across the two tasks. It also frames BioNN as a reusable testbed and suggests investigating these principles for large language models.

  • Homeostasis and structural plasticity most strongly accelerate the transition from memorization to generalization across both tasks.Homeostasis regulates long-term hidden-neuron use, while structural plasticity changes which connections participate in computation.
  • In sparse parity, the two mechanisms work together to form a compact, sparse effective circuit that generalizes rapidly.
  • In noisy XOR, homeostasis provides a robust route to generalization, while sufficiently sparse structural plasticity can induce generalization when homeostasis is absent.
  • Successful generalization is associated with reorganizing hidden representations through greater sparsity, specialization, and separation of latent structure.
  • BioNN preserves a common network while independently switching mechanisms to enable controlled comparisons of components and combinations.
  • For LLMs, the paper motivates testing whether homeostatic control and adaptive sparse computation can shorten representation reorganization and reduce optimization steps.
Loading 2608.28184v1…