Sustaining Plasticity via Learnable Wavelet Activations in Continual Learning

Sustaining Plasticity via Learnable Wavelet Activations in Continual Learning

Summary

This paper addresses the challenge of plasticity loss in continual learning, where models progressively lose the capacity to adapt to new tasks even when catastrophic forgetting is mitigated. The authors identify that standard fixed-form activation functions (e.g., ReLU) exhibit an inherent spectral bias towards low-frequency variations, which hinders the learning of high-frequency local details essential for complex tasks. Conversely, fully learnable activation paradigms (e.g., Rational approximations, KAN) allow unconstrained updates that can induce catastrophic forgetting at the activation level.

To address these limitations, the authors propose ChannelWavAct, a novel learnable wavelet-based activation framework that decomposes the activation function into a stable low-frequency base component and flexible high-frequency wavelet components. The activation for channel c is defined as:

Φc(xc) = w low,c · σ(xc) + Σ k=1 K w c,k · ψ((xc − τ c,k)/s k)

where σ(x) = SiLU(x) captures global low-frequency trends for stability, and ψ(u) = (u2 − 1)exp(−u2/2) (Mexican Hat wavelet) handles local high-frequency variations. The weights w c,k control amplitude, τ c,k shifts position, and s k controls the bandwidth (scale), with s k shared across channels to maintain consistent frequency structure.

The framework incorporates two key mechanisms:

1. Dynamic Wavelet Injection: A loss-based trigger mechanism monitors training dynamics using a stagnation counter C bad, which accumulates when the loss fails to drop below a relative margin δ from the minimum observed loss. When C bad exceeds a patience threshold P, the system injects ∆ new wavelet bases per channel. New bases are initialized with zero weights and translations, and scales based on the average of existing ones, to avoid abrupt perturbations.

2. Slope-Specific Regularization: A regularization term selectively penalizes deviations in pre-existing wavelet weights (w low and w c,k for k ≤ K old) while leaving geometric parameters (translations τ and scales s) free. This preserves core features from previous tasks while allowing adaptation to distribution shifts.

The authors also introduce decoupled optimization with two parameter groups: a plasticity group (backbone weights W and activation parameters w, τ, s) with an amplified learning rate η high = λ lr · η base, and a stability group (bias parameters b and batch normalization affine parameters) with the standard base rate η base. Additionally, post-activation batch normalization is applied to control variance shifts from the summation of multiple potentially unbounded functions.

Theoretical contributions include:

• Theorem 1: Under a uniform learning rate, high-frequency residual components decay negligibly due to the Dynamic Frequency Principle, motivating the need for decoupled learning rates.

• Theorem 2: The decoupled learning rate mechanism (η high > η base) mathematically compensates for the spectral bias, restoring plasticity for high-frequency information.

• Theorem 3: Optimization stagnation is characterized by the residual vector entering the approximate null space of the NTK, equivalent to Rayleigh quotient collapse. The minimum optimal incremental capacity is ∆* = ⌈µ(Ω res)/B ψ⌉, where µ(Ω res) is the Lebesgue measure of the residual's Fourier support and B ψ is the wavelet bandwidth.

• Theorem 4: The hybrid architecture (SiLU base + wavelets) is structurally necessary for efficient L2 approximation. A pure wavelet system requires divergent scales for non-zero-mean signals, while the SiLU base handles global trends and wavelets handle zero-mean residuals.

Experimental results demonstrate:

1. Trainability: On Permuted MNIST and Random Label MNIST (200 tasks each), ChannelWavAct achieves superior and stable accuracy compared to ReLU, AID, Randomized Smooth-Leaky, Rational, and B-Spline (KAN) baselines. It also maintains a high effective rank (srank) throughout training, unlike ReLU which suffers from rank collapse.

2. Generalizability:

• In replay-based SSD framework, ChannelWavAct achieves the highest Last and Average Accuracy on Mini-ImageNet (39.0% Avg), CIFAR100 (43.5% Avg), and Tiny-ImageNet (19.9% Avg).

• In replay-free EWC and WA frameworks, ChannelWavAct outperforms ReLU by 7.8% (CIFAR-100 EWC) and 22.6% (CIFAR-100 WA), demonstrating superior plasticity even under regularization constraints.

3. Spectral Analysis: Four experiments show ChannelWavAct achieves near-instantaneous convergence across all frequency components (5–30 Hz), maintains spectral integrity under parameter perturbations, resists low-frequency noise fitting, and decouples representational capacity from data manifold geometry.

4. Ablation Studies: Component analysis confirms that each element (learnable wavelet, dynamic injection, regularization) contributes to performance, with the full model achieving 43.5% Avg accuracy on CIFAR-100. Decoupled learning rates improve accuracy by 4.7% and reduce forgetting by 7.0%. Parameter sensitivity analysis shows optimal performance at ∆=1, P=50, λ=1, and δ=0.005.

The paper concludes that ChannelWavAct effectively harmonizes the stability-plasticity trade-off through its hybrid wavelet architecture, dynamic capacity expansion, and targeted regularization, achieving state-of-the-art performance across diverse continual learning benchmarks.

Improvements to AI Systems:

1. Adaptive Activation Functions with Frequency Decomposition

• Replace fixed activations (ReLU, GELU) with a hybrid SiLU + learnable wavelet basis per channel.

• The system can automatically separate stable low-frequency trends (via SiLU) from flexible high-frequency details (via Mexican Hat wavelets), enabling faster convergence on tasks with sharp local patterns (e.g., image edges, audio transients) without sacrificing global stability.

2. Dynamic Capacity Expansion Triggered by Optimization Stagnation

• Monitor loss reduction relative to a margin; when stagnation persists, inject new wavelet bases with zero-initialized weights.

• The system can autonomously grow its representational capacity only when needed, preventing over-parameterization while avoiding premature convergence to poor local minima—especially useful for non-stationary data streams.

3. Slope-Specific Regularization for Selective Memory Preservation

• Penalize only the amplitude weights of pre-existing wavelets, leaving translation and scale parameters free.

• The system can adapt to new tasks by shifting/reshaping existing features without overwriting the core learned representations, reducing catastrophic forgetting while still allowing fine-grained adjustments.

4. Decoupled Learning Rates for High-Frequency Plasticity

• Use amplified learning rates for backbone weights and activation parameters, while keeping bias and batch-norm parameters at a lower rate.

• The system can rapidly update high-frequency components (which are otherwise suppressed by spectral bias) while maintaining stable normalization statistics, enabling faster adaptation to distribution shifts without destabilizing training.

5. Post-Activation Batch Normalization for Unbounded Summations

• Apply batch norm after the summed activation output to control variance from multiple wavelet terms.

• The system can safely use complex, multi-term activations without risk of activation explosion, allowing deeper networks to maintain stable gradients over long training horizons.

6. Theoretically Grounded Capacity Sizing

• Use the derived formula ∆* = ⌈µ(Ω res)/B ψ⌉ to determine the minimum number of wavelets needed based on the residual’s Fourier support.

• The system can pre-estimate its required hidden-unit complexity for a given task, reducing trial-and-error in architecture design and improving sample efficiency.

---

What the Improved AI System Can Do:

• Continual Learning Without Forgetting: It can sequentially learn hundreds of tasks (e.g., permuted MNIST, random-label tasks) while maintaining high accuracy on both old and new tasks, even without replay buffers.

• Rapid Adaptation to New Domains: It can quickly fine-tune on new datasets (e.g., from CIFAR-100 to Tiny-ImageNet) with minimal performance drop on previously learned data, thanks to decoupled learning rates and selective regularization.

• Robustness to Label Noise and Distribution Shift: It can resist fitting low-frequency noise and maintain spectral integrity under perturbations, making it suitable for real-world data with noisy labels or covariate shift.

• Automatic Complexity Adjustment: It can dynamically add capacity only when training stagnates, avoiding wasteful computation while ensuring convergence on complex tasks—ideal for edge devices with limited memory.

• Interpretable Frequency Behavior: It can be analyzed via its wavelet coefficients to understand which frequency bands are being learned or forgotten, enabling debugging of model behavior in safety-critical applications (e.g., medical imaging, autonomous driving).