Thresholded Hidden-State Restart Gate

Integrate the gate into a GRU, LSTM, or diagonal state-space model immediately after computing the ordinary candidate state \(h_t^{\mathrm{raw}}=F_\theta(h_{t-1},u_t)\). Maintain an integer age \(t\), reset it to zero whenever the gate fires, and use \(h_t=(1-r_t)h_t^{\mathrm{raw}}+r_t h_{\mathrm{init}}\), where \(h_{\mathrm{init}}\) is either a learned vector or a small encoder of the current input. Compute a scalar deterioration score \(d_t=\operatorname{softplus}(w^\top[h_t^{\mathrm{raw}},u_t]-b)\). Start with a monotone threshold \(\tau(d)=\operatorname{clip}(\tau_0-\alpha d,1,T_{\max})\), with \(\alpha\geq0\), so larger estimated deterioration causes earlier restart. During training use the straight-through estimator: the forward gate is \(r_t=\mathbf1[t\geq\tau(d_t)]\), while the backward pass uses the sigmoid surrogate \(\tilde r_t=\sigma((t-\tau(d_t))/s)\). Add a restart penalty \(\lambda_r r_t\) and train with the task loss plus this penalty; anneal temperature \(s\) from 1 to 0.1. The paper's mathematical contribution is the constrained threshold policy and restart-versus-continuation Bellman comparison; estimate the deterioration score and threshold parameters empirically. First test on a 1-2 layer GRU and a diagonal SSM on WikiText-2 character prediction and Long Range Arena copy/addition tasks, comparing against the same model with its standard learned reset gate and fixed periodic resets. Measure validation loss, gradient norm, hidden-state norm drift, number of resets, and throughput. Success is lower long-context loss and fewer exploding or vanishing-gradient events at comparable reset count, or equal loss with substantially fewer recurrent state updates.