# 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.