chaobrain · xinzhu-L · Jul 20, 2026 · Jul 23, 2026
diff --git a/changelog.md b/changelog.md
diff --git a/docs/apis/nn/connectivity.rst b/docs/apis/nn/connectivity.rst
@@ -1,4 +1,4 @@
-Event-based Connectivity
+Event-based connectivity
 ========================
 
 .. currentmodule:: brainstate.nn

diff --git a/docs/apis/nn/conv.rst b/docs/apis/nn/conv.rst
@@ -1,4 +1,4 @@
-Convolutional Layers
+Convolutional layers
 ====================
 
 .. currentmodule:: brainstate.nn

diff --git a/docs/apis/nn/dropout.rst b/docs/apis/nn/dropout.rst
@@ -1,4 +1,4 @@
-Dropout Layers
+Dropout layers
 ==============
 
 .. currentmodule:: brainstate.nn

diff --git a/docs/apis/nn/dynamics.rst b/docs/apis/nn/dynamics.rst
@@ -1,4 +1,4 @@
-Dynamics and Simulation
+Dynamics and simulation
 =======================
 
 .. currentmodule:: brainstate.nn
@@ -7,7 +7,7 @@ Primitives for building dynamical systems and time-evolving neural models:
 differential-equation base classes, prefetch/delay utilities, ring-buffer delays,
 collective operations over module hierarchies, and numerical integration.
 
-Dynamics Base Class
+Dynamics base class
 -------------------
 
 ``Dynamics`` provides the foundation for differential equation-based models and
@@ -30,7 +30,7 @@ outputs are received during a simulation step.
    receive_update_input
    not_receive_update_input
 
-Dynamics Utilities
+Dynamics utilities
 ------------------
 
 Utilities for managing temporal dynamics, prefetching, and delayed outputs in
@@ -53,7 +53,7 @@ asynchronous signal processing in recurrent and spiking neural networks.
 
    init_maybe_prefetch
 
-Delay Utilities
+Delay utilities
 ---------------
 
 Temporal delay buffers and state management for neural dynamics with synaptic delays.
@@ -71,7 +71,7 @@ neural modeling. ``InterpolationRegistry`` manages interpolation methods for del
    StateWithDelay
    InterpolationRegistry
 
-Collective Operations
+Collective operations
 ---------------------
 
 Batch operations for managing states and function calls across module hierarchies.
@@ -92,7 +92,7 @@ and state management in complex neural architectures.
    vmap_reset_all_states
    assign_state_values
 
-Numerical Integration
+Numerical integration
 ---------------------
 
 Numerical integration methods for solving ordinary differential equations (ODEs)

diff --git a/docs/apis/nn/index.rst b/docs/apis/nn/index.rst
@@ -14,7 +14,6 @@ training utilities. Browse the reference by category below.
 
    module
    parameters
-   regularization
    linear
    conv
    normalization

diff --git a/docs/apis/nn/linear.rst b/docs/apis/nn/linear.rst
@@ -1,4 +1,4 @@
-Linear Layers
+Linear layers
 =============
 
 .. currentmodule:: brainstate.nn

diff --git a/docs/apis/nn/module.rst b/docs/apis/nn/module.rst
@@ -1,4 +1,4 @@
-Base Module Classes
+Base module classes
 ===================
 
 .. currentmodule:: brainstate.nn
@@ -17,7 +17,7 @@ of modules for feedforward architectures.
    ElementWiseBlock
    Sequential
 
-Common Wrappers
+Common wrappers
 ---------------
 
 Utility wrappers for context management and vectorization. ``EnvironContext`` manages

diff --git a/docs/apis/nn/normalization.rst b/docs/apis/nn/normalization.rst
@@ -1,4 +1,4 @@
-Normalization Layers
+Normalization layers
 ====================
 
 .. currentmodule:: brainstate.nn

diff --git a/docs/apis/nn/padding.rst b/docs/apis/nn/padding.rst
@@ -1,4 +1,4 @@
-Padding Layers
+Padding layers
 ==============
 
 .. currentmodule:: brainstate.nn

diff --git a/docs/apis/nn/parameters.rst b/docs/apis/nn/parameters.rst
@@ -1,4 +1,4 @@
-Parameter Containers
+Parameter containers
 ====================
 
 .. currentmodule:: brainstate.nn
@@ -16,7 +16,7 @@ Both support automatic caching of transformed values for improved performance.
    Param
    Const
 
-Parameter Transforms
+Parameter transforms
 --------------------
 
 Bijective transformations for constrained parameter optimization. These transforms
@@ -51,3 +51,60 @@ probabilistic applications. Use with ``Param`` for automatic constraint handling
    UnitVectorT
    ChainT
    MaskedT
+
+
+Standard regularizations
+------------------------
+
+Classical regularization methods for parameter penalization and constraint enforcement.
+These regularizations add penalty terms to the loss function to encourage desired
+properties like sparsity (L1), smoothness (L2), or structural constraints (orthogonality,
+spectral norms). Use with ``Param`` to automatically include regularization losses in
+training objectives.
+
+.. autosummary::
+   :toctree: ../generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   Regularization
+   L1Reg
+   L2Reg
+   ElasticNetReg
+   HuberReg
+   GroupLassoReg
+   TotalVariationReg
+   MaxNormReg
+   EntropyReg
+   OrthogonalReg
+   SpectralNormReg
+   ChainedReg
+
+Prior distribution-based regularizations
+-----------------------------------------
+
+Probabilistic regularizations based on prior distributions for Bayesian-inspired
+parameter estimation. These regularizations encode domain knowledge or assumptions
+about parameter distributions (Gaussian, heavy-tailed, bounded, etc.). Particularly
+useful for variational inference, maximum a posteriori (MAP) estimation, and
+uncertainty quantification. Each regularization implements ``loss()``, ``sample_init()``,
+and ``reset_value()`` for prior-based parameter initialization.
+
+.. autosummary::
+   :toctree: ../generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   GaussianReg
+   StudentTReg
+   CauchyReg
+   UniformReg
+   BetaReg
+   LogNormalReg
+   ExponentialReg
+   GammaReg
+   InverseGammaReg
+   LogUniformReg
+   HorseshoeReg
+   SpikeAndSlabReg
+   DirichletReg
diff --git a/docs/apis/nn/pooling.rst b/docs/apis/nn/pooling.rst
@@ -1,4 +1,4 @@
-Pooling and Reshaping
+Pooling and reshaping
 =====================
 
 .. currentmodule:: brainstate.nn

diff --git a/docs/apis/nn/regularization.rst b/docs/apis/nn/regularization.rst
diff --git a/docs/apis/nn/rnn.rst b/docs/apis/nn/rnn.rst
@@ -1,4 +1,4 @@
-Recurrent Cells
+Recurrent cells
 ===============
 
 .. currentmodule:: brainstate.nn

diff --git a/docs/apis/nn/utils.rst b/docs/apis/nn/utils.rst
@@ -1,4 +1,4 @@
-Utility Functions
+Utility functions
 =================
 
 .. currentmodule:: brainstate.nn

diff --git a/docs/apis/transform/autodiff.rst b/docs/apis/transform/autodiff.rst
@@ -1,4 +1,4 @@
-Gradient Computations
+Gradient computations
 =====================
 
 .. currentmodule:: brainstate.transform
@@ -8,7 +8,7 @@ Hessians. These functions extend JAX's autodiff capabilities with support for st
 computations, making them ideal for training neural networks and optimizing complex
 dynamical systems.
 
-Gradient Transformations
+Gradient transformations
 ------------------------
 
 .. autosummary::
@@ -18,7 +18,7 @@ Gradient Transformations
    vector_grad
    fwd_grad
 
-Vector-Jacobian and Jacobian-Vector Products
+Vector-Jacobian and Jacobian-Vector products
 --------------------------------------------
 
 .. autosummary::
@@ -38,7 +38,7 @@ Jacobian and Hessian
    jacobian
    hessian
 
-Base Classes
+Base classes
 ------------
 
 .. autosummary::

diff --git a/docs/apis/transform/control_flow.rst b/docs/apis/transform/control_flow.rst
@@ -1,4 +1,4 @@
-Control Flow
+Control flow
 ============
 
 .. currentmodule:: brainstate.transform

diff --git a/docs/apis/transform/core.rst b/docs/apis/transform/core.rst
@@ -1,4 +1,4 @@
-Core Transformation Infrastructure
+Core transformation infrastructure
 ==================================
 
 .. currentmodule:: brainstate.transform

diff --git a/docs/apis/transform/debugging.rst b/docs/apis/transform/debugging.rst
@@ -1,4 +1,4 @@
-Debugging and Error Checking
+Debugging and error checking
 ============================
 
 .. currentmodule:: brainstate.transform

diff --git a/docs/apis/transform/ir.rst b/docs/apis/transform/ir.rst
@@ -1,4 +1,4 @@
-Intermediate Representation (IR) Tooling
+Intermediate Representation (IR) tooling
 ========================================
 
 .. currentmodule:: brainstate.transform
@@ -7,7 +7,7 @@ Tools for optimizing, processing, generating code from, and visualizing JAX
 intermediate representations (Jaxpr). These utilities reduce computation overhead and
 improve runtime performance while preserving a function's semantics and interface.
 
-IR Optimization
+IR optimization
 ---------------
 
 Optimize Jaxpr intermediate representations by applying compiler optimizations such as
@@ -25,7 +25,7 @@ propagation, and algebraic simplification.
    algebraic_simplification
    optimize_jaxpr
 
-IR Processing and Transformation
+IR processing and transformation
 --------------------------------
 
 Tools for processing and transforming JAX intermediate representations, including
@@ -39,7 +39,7 @@ equation-to-Jaxpr conversion and JIT inlining operations.
    eqns_to_jaxpr
    inline_jit
 
-Code Generation
+Code generation
 ---------------
 
 Convert JAX functions and Jaxpr representations into readable Python code for

diff --git a/docs/apis/transform/jit.rst b/docs/apis/transform/jit.rst
@@ -1,4 +1,4 @@
-JIT Compilation
+JIT compilation
 ===============
 
 .. currentmodule:: brainstate.transform

diff --git a/docs/apis/transform/mapping.rst b/docs/apis/transform/mapping.rst
@@ -1,4 +1,4 @@
-Mapping and Vectorization
+Mapping and vectorization
 =========================
 
 .. currentmodule:: brainstate.transform
@@ -7,7 +7,7 @@ Transformations for vectorized and parallel computation across multiple data poi
 or devices. These functions enable efficient batch processing and multi-device
 scaling, essential for large-scale simulations and distributed training.
 
-Basic Vectorization
+Basic vectorization
 -------------------
 
 Vectorize computations across batch dimensions. ``vmap2`` is the recommended API
@@ -22,7 +22,7 @@ with enhanced state handling and control over batching axes.
    vmap2_new_states
    map
 
-Parallel Mapping
+Parallel mapping
 ---------------
 
 Execute computations in parallel across devices, or shard them with explicit mesh
@@ -35,7 +35,7 @@ control.
    pmap2_new_states
    shard_map
 
-Base Classes and Utilities
+Base classes and utilities
 --------------------------
 
 .. autosummary::

diff --git a/docs/concepts/the_graph_model.md b/docs/concepts/the_graph_model.md
@@ -1,4 +1,4 @@
-# The Graph Model
+# The graph model
 
 A BrainState model is an ordinary Python object graph: modules hold sub-modules in attributes,
 lists, and dicts, and {class}`~brainstate.State` objects sit at the leaves. The graph model is the
-Original file line number
+Diff line change
@@ Expand Up / @@ -14,7 +14,6 @@ training utilities. Browse the reference by category below. @@
        module
        parameters
-       regularization
        linear
        conv
        normalization
@@ Expand Down @@