code/Taotern_LLM_Experiments/docs/CURRENT_SSM_LLM_ARCHITECTURE.md

Current SSM LLM Architecture

This document explains the current TaoNet-SSM and TaoNet-Hybrid models from the LLM surface down to the DPLR SSM matrices. It describes the SSM implementation and the current best real-token candidate in this experiment ledger.

Current Best Candidate

The current best real-token candidate is now the hybrid TaoNet:

architecture_type = taonet_hybrid
block order       = SSM, attention, SSM, attention
hidden_dim        = 256
num_layers        = 4
num_heads         = 4
hidden_dim_ff     = 1024
ssm_core          = dplr
ssm_hidden_dim    = 16
ssm_mixer_dim     = 128
ssm_rank          = 1
ssm_kernel_mode   = conv
ssm_local_shift   = true
shift gain        = per-channel
finite tail       = disabled for the current best speed/quality point
SSM gate type     = channel
SSM lanes         = 2
lane mode          = split
split mix         = none for current batch64 best; Hadamard helps only batch32 hybrid
lane combine      = concatenation after split lanes
dtype             = bf16
current SSM commit = 76f725f
current TaoTrain commit = 89aa98d

The current best pure SSM candidate is:

architecture_type = taonet_ssm
hidden_dim        = 256
num_layers        = 4
num_heads         = 4
hidden_dim_ff     = 1024
ssm_core          = dplr
ssm_hidden_dim    = 16
ssm_mixer_dim     = 128
ssm_rank          = 1
ssm_kernel_mode   = conv
ssm_local_shift   = true
shift gain        = per-channel
finite tail       = disabled
SSM gate type     = channel
SSM lanes         = 2
lane mode          = full currently has best pure-SSM quality; split is the faster/lower-memory candidate
split mix         = Hadamard tested but not better overall for pure SSM
lane combine      = channel for full lanes, concatenation for split lanes
dtype             = bf16
current SSM commit = 76f725f
current TaoTrain commit = 89aa98d

The attention baseline is taonet with the same outer dimensions. The point of the comparison is to keep the LLM scaffold nearly fixed and replace only the sequence-mixing core. The DPLR direct path supports an exact finite-response readout rewrite:

C @ (response - z^L A^L response)
  = C @ response - z^L (C @ A^L) @ response

This improves forward-only long-context timing but has not improved the training-time forward+backward path enough to become the current benchmark default. The current best benchmark uses the faster approximate finite-tail path.

The latest large run showed that fully replacing attention with SSM is not the best path at the current scale, but combining attention and SSM is promising. taonet_hybrid alternates the original TaoNet attention blocks and the SSM blocks while sharing the same embedding, final norm, output head, model dimension, FFN width, and training script.

For four layers, the current best hybrid computes:

x_0 = token_embedding(tokens)
x_1 = SSM_TaoNet_Block_1(x_0)
x_2 = Attention_TaoNet_Block_2(x_1)
x_3 = SSM_TaoNet_Block_3(x_2)
x_4 = Attention_TaoNet_Block_4(x_3)
h   = final_layer_norm(x_4)
logits = output_head(h)

Top-Level LLM

The model is an autoregressive next-token language model.

Given token ids:

tokens: [batch, seq]

the model computes:

x_0 = token_embedding(tokens)          # [batch, seq, d_model]
x_0 = dropout(x_0)

for layer l in 1..L:
    x_l = SSM_TaoNet_Block_l(x_{l-1})

h = final_layer_norm(x_L)
logits = output_head(h)                # [batch, seq, vocab]
loss = cross_entropy(logits[:, t], labels[:, t])

In the current pilot:

vocab = 8192
d_model = 256
L = 4

The output head is a dense linear projection from d_model to the tokenizer vocabulary.

Block Structure

Each SSM TaoNet block preserves the normal Transformer-like residual layout:

x' = x + dropout(SSM_Mixer(LN_1(x)))
y  = x' + dropout(SwiGLU_FFN(LN_2(x')))

The feed-forward branch is unchanged from TaoNet-style dense FFN:

g = W_gate LN_2(x')
v = W_value LN_2(x')
ff = W_out (SiLU(g) * v)

For current dimensions:

W_gate  : 256 -> 1024
W_value : 256 -> 1024
W_out   : 1024 -> 256

This means most per-block parameters still sit in the dense FFN, not in the small DPLR SSM core.

SSM Mixer Structure

The SSM mixer receives:

x: [batch, seq, d_model]

where d_model = 256.

The current mixer path is:

x_norm = LayerNorm(x)

if input_gate:
    x_gated = x_norm * sigmoid(W_input_gate x_norm + b_input_gate)
else:
    x_gated = x_norm

u = W_in x_gated                    # [batch, seq, ssm_mixer_dim]

if ssm_lane_mode == split:
    lane_inputs = split(u, ssm_num_lanes, dim=-1)
else:
    lane_inputs = [u, ..., u]

for lane i in 1..ssm_num_lanes:
    y_i = DPLR_SSM_i(lane_inputs_i)

if ssm_lane_mode == split:
    if ssm_split_mix == hadamard:
        y_ssm = concat(y_1 + y_2, y_1 - y_2) / sqrt(2)
    else:
        y_ssm = concat_i(y_i)
elif ssm_num_lanes == 1:
    y_ssm = y_1
elif lane_combine == mean:
    y_ssm = mean_i(y_i)
elif lane_combine == channel:
    y_ssm[:, :, c] = sum_i lane_weight[i, c] * y_i[:, :, c]

y_ssm = activation(y_ssm)
y_proj = W_out y_ssm                # [batch, seq, d_model]

if output_gate:
    y_proj = y_proj * sigmoid(W_output_gate x_norm + b_output_gate)

y_proj = layer_scale * y_proj

if local_shift:
    y_proj[:, t] += shift_scale * x_norm[:, t-1]

return dropout(y_proj)

Current settings:

d_model       = 256
ssm_mixer_dim = 128
ssm_num_lanes = 2 for the current best hybrid quality point
ssm_lane_mode = split for the current best hybrid efficiency-quality point
ssm_split_mix = none for the current batch64 best
input_gate    = enabled
output_gate   = enabled
activation    = GELU
layer_scale   = learned vector of length 256
local_shift   = enabled
shift_scale   = learned vector of length 256

The local shift branch is deliberately simple:

shifted_t = x_norm_{t-1}
output_t += alpha * shifted_t

where alpha is per-channel in the current best setting. This branch was the major quality breakthrough for previous-token memory and also improves real-token modeling.

The multi-lane branch is the latest SSM-capacity improvement. In full-lane mode, each lane has independent DPLR parameters and sees the same projected input u. The channel combine is:

y_t,c = sum_i w_i,c * y_i,t,c

with:

w : [num_lanes, ssm_mixer_dim]

This combine operation is elementwise across channels, so it is friendly to ternary deployment.

In split-lane mode, the projected channels are partitioned before the DPLR core:

u = [u_1, u_2, ..., u_L]
y_i = DPLR_SSM_i(u_i)
y = concat_i(y_i)

For the current ssm_mixer_dim=128 and ssm_num_lanes=2, each split lane has 64 channels. The high-scale benchmark shows split lanes recover much of the throughput and memory lost by full-lane duplication, and the best hybrid now uses split lanes. Pure SSM still trails attention, so the next planned version should add cheap cross-channel communication after the split-lane SSM output.

The fixed Hadamard split mix was tested as:

y = concat(y_1 + y_2, y_1 - y_2) / sqrt(2)

It adds no learned parameters and is ternary-friendly, but the high-scale run showed it is too rigid: it improves the batch-32 ssm_first hybrid slightly, does not improve pure SSM enough, and does not beat plain split lanes at batch 64. The next cross-channel mix should be learnable while remaining ternary-friendly.

DPLR SSM Core

The SSM core operates in the projected mixer dimension:

u_t in R^m
h_t in R^n
y_t in R^m

where the current best candidate uses:

m = ssm_mixer_dim  = 128
n = ssm_hidden_dim = 16
r = ssm_rank       = 1

The recurrent form is:

h_t = A_bar h_{t-1} + B_bar u_t
y_t = C h_t + D * u_t

Matrix shapes:

A_bar : [n, n]
B_bar : [n, m]
C     : [m, n]
D     : [m]

The current DPLR structure defines:

A_bar = diag(a) - U V^T

with:

a : [n]
U : [n, r]
V : [n, r]

For the current best candidate:

A_bar : [16, 16]
B_bar : [16, 128]
C     : [128, 16]
D     : [128]
U     : [16, 1]
V     : [16, 1]

Continuous-to-Discrete Parameters

The diagonal continuous eigenvalues are stable by construction:

lambda_i = -softplus(log_lambda_real_i)

The learned step size is:

dt = clamp(softplus(log_dt), dt_min, dt_max) * rate

The diagonal discrete transition is:

a_i = exp(dt * lambda_i)

The input matrix is discretized per state:

B_bar_i = ((a_i - 1) / lambda_i) * B_i

with the small-lambda fallback:

B_bar_i = dt * B_i

The low-rank factors are ternary-aware:

U = ternary_u_mask * softplus(log_u_amp) * max_low_rank_scale * sigmoid(low_rank_logit)
V = ternary_v_mask * softplus(log_v_amp)

The masks are fixed buffers with entries in:

{-1, 0, 1}

The amplitudes are learned real values. This is why the current DPLR core is best described as ternary-aware rather than fully ternary. The sign structure is ternary, while amplitudes, projections, gates, FFN, and embeddings are still dense learned tensors.

Frequency-Domain Training Path

For sequence training, the model uses the convolutional path instead of stepping token by token.

The recurrent SSM implies a convolution kernel:

K_k = C A_bar^k B_bar

The output can be computed as:

y_t = sum_{k=0..t} K_k u_{t-k} + D * u_t

The implementation uses FFT:

U_f = rfft(u)
Y_f = H(z) U_f
y   = irfft(Y_f)

The transfer response is computed with the DPLR inverse:

H(z) = C (I - z^L A_bar^L) (I - z A_bar)^(-1) B_bar + D

The inverse is not formed as a dense inverse in the main direct path. Instead, it uses Woodbury-style DPLR algebra:

(I - z (diag(a) - U V^T))^-1

which reduces the low-rank correction to small rank operations. Since the current best uses rank 1, the correction is especially small.

Parameter Inventory For h16/m128

Inside one SSM mixer block, the DPLR core has approximately:

Parameter	Shape	Count
log_lambda_real	[16]	16
B	[16, 128]	2048
C	[128, 16]	2048
D	[128]	128
log_u_amp	[1, 16]	16
log_v_amp	[1, 16]	16
low_rank_logit	[1]	1
log_dt	scalar	1
DPLR total		4274

The mixer around the SSM has larger dense components:

Component	Shape	Count
Input projection W_in	[256, 128]	32768
Output projection W_out	[128, 256]	32768
Input gate	256 -> 256 with bias	65792
Output gate	256 -> 256 with bias	65792
Layer scale	[256]	256
Per-channel local shift	[256]	256

The FFN is larger again:

Component	Shape	Count
FF gate	256 -> 1024	262144
FF value	256 -> 1024	262144
FF output	1024 -> 256	262144
FFN total		786432

So, in the current LLM, the SSM core is small. Performance and deployment friendliness depend not only on the DPLR kernel, but also on projections, gates, local shift, FFN, embeddings, and output head.

Inference Export

The DPLR core exposes inference matrices:

A_continuous_diag
A_discrete
low_rank_U
low_rank_V
B
C
D
dt
ternary_u_mask
ternary_v_mask
diag

The step-wise inference recurrence is:

h_new = h A_bar^T + u B_bar^T
y     = h_new C^T + u * D

This is the path to target for ternary deployment analysis.

Current Evidence Against Attention TaoNet

Latest high-scale comparison with attention, pure SSM, and hybrid on /home/student/Data/TaoData/pretrain.jsonl:

Batch	Model	Pattern	Gate	Eval loss	Eval accuracy	Forward+backward tok/s
32	attention TaoNet	-	-	3.5270	0.3490	1.367M
32	SSM TaoNet h16/m128	-	dense	3.7575	0.3138	1.144M
32	SSM TaoNet h16/m128	-	channel	3.7708	0.3109	1.158M
32	hybrid TaoNet	ssm_first	channel	3.4733	0.3563	1.239M
64	attention TaoNet	-	-	3.3949	0.3647	1.449M
64	SSM TaoNet h16/m128	-	dense	3.6478	0.3257	1.254M
64	SSM TaoNet h16/m128	-	channel	3.6554	0.3252	1.231M
64	hybrid TaoNet	ssm_first	channel	3.3694	0.3693	1.325M

Interpretation:

• Hybrid is now the best candidate in the ledger: the channel-gated ssm_first pattern beats attention on eval loss and token accuracy at both tested batch sizes.
• Hybrid is still slower than attention, but it keeps about 91% of attention forward+backward throughput.
• Pure SSM with either dense or channel gates did not beat attention in the high-scale run.
• Channel gates are more ternary-friendly and improved the best hybrid, but slightly hurt pure SSM quality.
• The next pure-SSM question is how to add SSM capacity without falling back to attention-like dense mixing.

Latest optimizer sweep for the current h16/m128 candidate:

Model	LR	Eval loss	Eval accuracy	Forward+backward tok/s
attention TaoNet	0.0008	4.688	0.215	1.37M
SSM TaoNet h16/m128	0.0004	4.952	0.207	1.09M
SSM TaoNet h16/m128	0.0006	4.789	0.218	1.08M
SSM TaoNet h16/m128	0.0008	4.716	0.224	1.08M
SSM TaoNet h16/m128	0.0012	4.705	0.223	1.09M

Interpretation:

• Attention still has the best validation loss, but the gap is still small: latest attention 4.688 vs best SSM 4.705.
• SSM has the better token accuracy at useful learning rates: 0.224 at LR 0.0008 and 0.223 at LR 0.0012.
• SSM throughput is currently behind attention in the controlled optimizer run.
• The follow-up weight-decay sweep at LR 0.0012 did not improve quality beyond wd=0.01.
• The DPLR discrete-parameter reuse cleanup preserved quality and gave only a small speed movement at the TaoNet level.
• A rank-one frequency specialization was attempted and reverted after it regressed SSM forward+backward throughput to about 497k tok/s.
• Component profiling at the TaoData benchmark shape shows the DPLR SSM core is the largest measured SSM-side forward cost: about 2.203 ms/forward across four layers. Gates are about 0.403 ms/forward, projections about 0.303 ms/forward, FFN linears about 0.875 ms/forward, and the output head about 0.384 ms/forward.
• DPLR microprofiling shows the direct frequency path is better than materialized transfer at the current shape: 1.812 ms vs 2.478 ms forward+backward, and 119 MB vs 308 MB peak allocation.
• A batch-major direct-path layout rewrite was attempted and reverted after it regressed forward+backward to 3.435 ms.
• The next unresolved question is whether direct-path copy/cast and complex batched-matmul overhead can be reduced with a lower-level kernel plan without losing the accuracy edge.

Known Limitations

• The current SSM LLM is ternary-aware, not fully ternary.
• Dense projections, gates, FFN, embeddings, and output head still dominate parameter count.
• Batch generalization for h16/m128 has been checked at batch 16, 32, and 64; SSM kept its accuracy edge but not a loss or speed lead.
• The frequency-domain DPLR path still depends on complex FFT and complex rank-one algebra.
• Hardware acceleration work should target the DPLR frequency response/backward path and the inference recurrence separately.