Hard-Constrained Bernstein Memory Head

(1) Integration point: insert the module after the value projection in each Transformer block, or beside the state update in an SSM. For hidden states $h_{1:T}\in\mathbb{R}^{T\times d}$, compute values $v_t=W_vh_t$, construct the causal memory vector $m_t$, and add it through a learned projection and scalar gate as $h_t^{\mathrm{out}}=h_t+gW_mm_t$. Start with one shared kernel per layer; optionally use one kernel per attention head.

(2) Pseudocode:

input h[1:T]
v = h @ Wv
raw = softplus(theta[m+1])
q = raw
for j = m-1 down to 0:
    d[j] = sum(q[j:m])
c[m] = softplus(a)
for j = m-1 down to 0:
    c[j] = c[j+1] + d[j]
for z in lag_grid:
    k[z] = sum_j c[j] * binom(m,j) * z**j * (1-z)**(m-j)
w = k(lag_grid) / (sum(k(lag_grid)) + eps)
for t in 1..T:
    memory[t] = sum_{r=1}^{min(L,t-1)} w[r] * v[t-r]
out = h + sigmoid(g) * (memory @ Wm)

The nested cumulative sums enforce positivity, monotone decrease, and convexity without penalty losses. Use torch.flip and cumulative sums to vectorize coefficient construction, and use a causal depthwise convolution or FFT convolution for long contexts.

(3) Computed from the paper's mathematics: the Bernstein basis, coefficient-difference constraints, and the causal kernel convolution. Estimated empirically: the degree $m$, memory horizon $\tau$, gate $g$, projection matrices, and whether one shared kernel or per-head kernels are preferable. If $\tau$ is learned, parameterize it as softplus(tau_raw) and clip it to a practical range. Monitor the discrete first differences $k_r-k_{r+1}$ and second differences $k_r-2k_{r+1}+k_{r+2}$ to verify the implementation.

(4) First cheap experiment: train a 4-layer, 8-head Transformer on synthetic autoregressive sequences generated by a known mixture of exponential and power-law memory, with random input corruption and only 10-20% observed training tokens. Compare the constrained head against standard attention, an unconstrained learned depthwise lag convolution with the same number of parameters, and a single exponential kernel. Evaluate validation MSE, stability on sequences 4x longer than training, and sensitivity to noise. Success is lower extrapolation error and fewer exploding or oscillating runs at equal FLOPs, while the learned kernel shape remains stable across random seeds.