Localized Spectral Redistribution

Integrate this at the hidden-state propagation operation of a residual SSM, RNN, or graph neural network, rather than at the optimizer. Let h be the current hidden tensor and let A(h) denote the ordinary one-step propagation. Construct a sparse local redistribution map R: for a sequence, a flagged token or channel is replaced by a weighted average of itself and its preceding token; for a graph, a flagged node is replaced by a degree- or attention-weighted average over its one-hop neighbors. Every K training steps, estimate local amplification. Use JVPs to compute v=J_A(h)u for a random probe u, then estimate each coordinate's gain L_i=|v_i|/(|u_i|+\epsilon); for a more robust estimate, run two or three power iterations on each node or channel block. Flag coordinates with L_i>2. For every flagged coordinate set s_i=clip((L_i-2)/(L_i-1)+m,0,1), where m is a safety margin such as 0.05, and compute h_new=(1-s)\odot A(h)+s\odot R(A(h)); unflagged coordinates use s_i=0. For a sequence, use R(h)_t=(w_0h_t+w_1h_{t-1})/(w_0+w_1), with w_0 and w_1 determined by normal and flagged coordinate masses. For graphs, use R(h)_i=\sum_{j\in N(i)\cup\{i\}}w_{ij}h_j/\sum_jw_{ij}. Stop gradients through the gain estimate and s. The threshold formula is taken from the paper; Jacobian gains, neighbor weights, and flagged coordinates are estimated empirically. First test a 4-8 layer linearized SSM and a 3-layer message-passing GNN on synthetic long sequences or Cora. Compare ordinary propagation, global residual damping, and selective redistribution at equal FLOPs. Record maximum Jacobian singular value, exploding-gradient frequency, validation accuracy, and the largest usable step size or depth. Success means fewer NaNs and a lower spectral radius at equal accuracy, or equal stability with less global damping and faster loss descent.