Skip to content

v-code01/stabgauge

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

stabgauge

The stabilizer that keeps xLSTM's sLSTM from overflowing is a gauge on its state: the output does not depend on the stabilizer offset at all. The deployed running-max offset is the pointwise minimal overflow-safe gauge, and it converts the normalizer's exponential blow-up, which overflows a double at an exact step, into at-most-linear growth.

xLSTM's sLSTM (Beck et al. 2024) carries a cell state and a normalizer with exponential input and forget gates and returns their ratio:

C_t = f_t C_{t-1} + i_t z_t ,   N_t = f_t N_{t-1} + i_t ,   h_t = C_t / N_t ,

with i_t = exp(logi_t), f_t = exp(logf_t). Because the exponential gates overflow, the layer carries a stabilizer offset m_t and shifts the gates to f'_t = exp(logf_t + m_{t-1} - m_t) and i'_t = exp(logi_t - m_t).

The stabilizer is a gauge

For any offset sequence, both C_t and N_t acquire the common factor exp(-m_t) by induction, so their ratio h_t is exactly invariant to m_t. The stabilizer is a gauge transformation on the state, and the running max is only one gauge choice: the naive (m_t = 0), the running-max, and an arbitrary random offset schedule all produce the same output (measured |h_naive - h_max| and |h_random - h_max| at machine zero, 3e-16 and 3e-14). This is the same gauge idea as the softmax and cross-entropy logit shifts, here on the sLSTM cell and normalizer.

The running max is the minimal safe gauge

Among the offsets that keep both shifted gates at most 1, and so avoid overflow, the running max m_t = max(logi_t, logf_t + m_{t-1}) is the pointwise smallest. It attains max(i'_t, f'_t) = 1 exactly at every step (measured 1200/1200), and lowering it by any epsilon pushes a gate strictly above 1 toward overflow (measured 1200/1200). So the deployed stabilizer is the tight, minimal overflow-safe gauge, not just a safe one.

The exact overflow horizon

With constant log-gates a = logf, b = logi, the unstabilized normalizer is a geometric sum N_t = e^b (f^t - 1)/(f - 1), f = e^a, growing exponentially, so it exceeds the largest finite double at the exact horizon

t* = floor( (L - b + ln(f - 1)) / a ) + 1 ,   L = ln(DBL_MAX) = 709.78 ,

carrying a -b + ln(f - 1) correction from the normalizer geometry (measured overflow at t = 1415, formula t* = 1415, matched to the step). Under any overflow-safe gauge both gates are at most 1, so N_t <= t and |C_t| <= t max|z|: the stabilizer turns exponential blow-up into at-most-linear growth (measured max N = 2.54 out to t = 200000), overflow-free for every realistic length, while keeping the output bit-invariant.

Layout

  • slstm.py: the sLSTM scalar recurrence run under any stabilizer offset schedule.
  • gauge.py: the output invariance to the offset (gauge freedom) and the running-max minimality (tightness).
  • overflow.py: the exact t* overflow horizon of the unstabilized normalizer and the linear boundedness of the stabilized form.
  • test_stabgauge.py: the gauge invariance to the max and to an arbitrary offset, the minimality, the horizon formula, the linear bound, and the steady state.

Reproduce

python slstm.py
python gauge.py
python overflow.py
python test_stabgauge.py

Note

The xLSTM paper introduces the running-max stabilizer as a safe-softmax-style fix, stating that it keeps the exponent argument at most 1 and so bounds the gates. That gate-boundedness is not claimed here. What is here is the sharper structure: the stabilizer offset is a gauge, so the output is invariant to any offset sequence and not only the max; the running max is the pointwise minimal overflow-safe gauge, attaining max(gate) = 1 tightly; the unstabilized normalizer overflows at the exact horizon t* = floor((L - b + ln(f - 1))/a) + 1 with its normalizer-geometry correction; and the stabilizer converts exponential growth to linear while leaving the output bit-invariant. It is distinct from the shipped scanhorizon (an SSM associative-scan divide-out-gate horizon with no normalizer or gauge), covfirewall (the mLSTM readout clamp, a cancellation rather than overflow question), and gatebias (a bf16 compounding result). The named extension is the time-varying-gate horizon as a function of the log-gate path and a bf16 measurement of where the linear-bounded form still loses precision.

About

xLSTM sLSTM's overflow stabilizer is a GAUGE on the (C,N) state: output h=C/N is exactly invariant to any offset, not just the running max. The max is the pointwise-MINIMAL overflow-safe gauge (hits max(gate)=1 tight); the unstabilized normalizer overflows at exact t*=floor((L-b+ln(f-1))/a)+1, stabilized grows only linearly. Proof + fp64.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages