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:

```text

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:

```text

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:

```text

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:

```text

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:

```text

tokens: [batch, seq]

```

the model computes:

```text

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:

```text

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:

```text

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:

```text

g = W_gate LN_2(x')

v = W_value LN_2(x')

ff = W_out (SiLU(g) * v)

```

For current dimensions:

```text

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:

```text

x: [batch, seq, d_model]

```

where `d_model = 256`.

The current mixer path is:

```text

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:

```text

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:

```text

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:

```text

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

```

with:

```text

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:

```text

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:

```text

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:

```text

u_t in R^m

h_t in R^n

y_t in R^m

```

where the current best candidate uses:

```text

m = ssm_mixer_dim  = 128

n = ssm_hidden_dim = 16

r = ssm_rank       = 1

```

The recurrent form is:

```text

h_t = A_bar h_{t-1} + B_bar u_t

y_t = C h_t + D * u_t

```

Matrix shapes:

```text

A_bar : [n, n]

B_bar : [n, m]

C     : [m, n]

D     : [m]

```

The current DPLR structure defines:

```text

A_bar = diag(a) - U V^T

```

with:

```text

a : [n]

U : [n, r]

V : [n, r]

```

For the current best candidate:

```text

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:

```text

lambda_i = -softplus(log_lambda_real_i)

```

The learned step size is:

```text

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

```

The diagonal discrete transition is:

```text

a_i = exp(dt * lambda_i)

```

The input matrix is discretized per state:

```text

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

```

with the small-lambda fallback:

```text

B_bar_i = dt * B_i

```

The low-rank factors are ternary-aware:

```text

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:

```text

{-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:

```text

K_k = C A_bar^k B_bar

```

The output can be computed as:

```text

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

```

The implementation uses FFT:

```text

U_f = rfft(u)

Y_f = H(z) U_f

y   = irfft(Y_f)

```

The transfer response is computed with the DPLR inverse:

```text

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:

```text

(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:

```text

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:

```text

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.