Diffusion Router Training and Research Formula DRAFT

diffusion-router-training-draft.md

Diffusion Router Training — draft algorithm, two alternatives

Date: 2026-09-09 · Status: draft (research, not yet scheduled) Companion: docs/2026-09-09-bandwidth-model.md, docs/2026-09-09-ftq-healing-plan-v3.md, ~/labware/llocollama/plans/GemFlash38-on_algebra.md (Log-Exp theory), ~/labware/llocollama/plans/GemFlash38_om_superfluid_bi-directional_attention.md (counterflow theory), ~/labware/llocollama/plans/GemFlash38_superfluid_state_fix_DEFERRED.md (deferred state fix)

The research question

Can the 512-expert MoE layers of Qwen3.8-Flash-Next operate under bidirectional attention (diffusion-canvas denoising) instead of causal autoregressive attention — and if so, does the batched expert streaming (the bandwidth model's §5) yield the predicted 5× streaming amortization on memory-constrained hardware?

The expert MLPs are architecturally agnostic to the attention mask (position-independent per-token MLPs). The routing is the coupling point: the router (ffn_gate_inp) sees the hidden state produced by attention, so bidirectional context changes which experts fire. Fine-tuning under the diffusion objective lets the router and experts adapt to the bidirectional representation.

Alternative A: SOTA / stock diffusion-MoE fine-tune

The standard approach, as done for DiffusionGemma (Gemma-2 → block diffusion) and DeepSeek-V4-Flash-style MoE models. Nothing custom; every component uses established methods.

Model surgery (one-time, at training start)

Step	What	Why
A.1	Remove causal mask from QSA attention layers (12 of 48; every 4th layer)	enables bidirectional context on full-attention layers
A.2	GDN layers: run the recurrence bidirectionally (forward scan → h_fwd, backward scan → h_bwd, concatenate h_fwd⊕h_bwd)	linear attention layers are inherently causal; the two-pass combine is the standard SOTA for bidirectional linear attention (Mamba-2, RWKV-v6 variants)
A.3	Keep MoE structure, SiLU gating, HC mixer, PLE injection unchanged	expert MLPs are position-independent; the router adapts via fine-tuning
A.4	Keep the PLE/engram injection but remap the n-gram hash: "committed tokens" → "currently-visible (unmasked) canvas tokens"	the PLE is token-keyed, not position-keyed; the n-gram hash naturally operates on the unmasked subset

Training loop

for step in 1..N_steps (300–500):
    batch = sample_human_text(batch_size=8, seq_len=2048)
    canvas = tokenize(batch)                                # [8, 2048]

    # corruption schedule (E1/E2 block-diffusion)
    mask_ratio = cosine_schedule(step / N_steps)            # 0.1 → 0.5
    mask = random_block_mask(canvas, mask_ratio)             # contiguous blocks

    # forward (bidirectional)
    R = embed(unmasked_tokens, mask_tokens)                  # [8, 2048, 4×2560] HC streams
    R = inject_ple(R, ngram_hash(visible_tokens))            # PLE engram lookup + gate
    for layer in 0..47:
        x = hc_mix_attn(R)
        y = (layer % 4 == 0) ? QSA_bidir(x) : GDN_bidir(x)  # bidirectional hybrid
        R = hc_combine_attn(R, y)
        x = hc_mix_ffn(R)
        routes = router(x)                                   # [8, 2048, 384]
        y = Σ_{e ∈ top-10(routes)} gate_e·down_e·silu(up_e·x)   # MoE with SiLU
        R = hc_combine_ffn(R, y)
    logits = lm_head(final_hc_norm(R))                       # [8, 2048, 248320]

    # loss: CE on masked positions only
    loss = cross_entropy(logits[mask], target_tokens[mask])
    + λ_anchor · KL(logitsFTQ || logits_original)             # anchor to the causal
                                                              # model's distribution
                                                              # on unmasked positions
    loss.backward()
    optimizer.step()                                          # warm tier + router only

Training parameters

Parameter	Value	Rationale
Steps	300–500	router-adaptation + expert-function recovery (referee estimate)
Learning rate	1e-4 (warm tier), 1e-5 (router)	router is sensitive (referee correction #2)
Trainable	warm experts (0.42B) + router (0.39M)	hot tier frozen, backbone frozen
Optimizer	8-bit Adam (bitsandbytes 0.50.2)	1.7 GB state on 0.42B params
λ_anchor	0.1–0.5	prevents catastrophic drift from the causal distribution
Batch	8 × 2048 tokens	canvas size = context size
Corruption	block mask, cosine schedule	DiffusionGemma E1/E2 schedule

Compute

Phase	Where	Wall-clock
Canvas training	128 GB rental (student fp8 + warm bf16)	12–24 h at 300–500 steps
PLE re-derivation	after bidirectional fine-tune, the n-gram hash inputs change	~1 h (re-index)
Validation	ppl + 35B-A3B benchmark	on the 2070S

Expected outcome

The router adapts to the bidirectional representation; the warm experts recover the folded-expert function under the new attention. The expert streaming amortization (batched canvas MoE) yields the bandwidth model's predicted ~5× per-token streaming reduction. Quality = the 35B-A3B benchmark.

SOTA risks (referee-flagged)

• Router collapse under bidirectional context (the flat-routing histogram means the router has no strong prior to fall back on). Mitigation: router warmup phase (referee correction #2).
• GDN bidirectional two-pass doubles the GDN state memory: 48 GDN layers × state size — needs gradient checkpointing on the GDN backward pass.
• Catastrophic drift: the causal→bidirectional transition is a large distribution shift; the KL anchor to the original model's unmasked-position logits is essential.

---

Alternative B: Log-Exp Pentagonator (research formula)

Replaces three components with the research formula:

1. Expert gating: SiLU → Log-Exp pentagonator (the algebra doc's Hopfield + Lorentzian shape)
2. GDN bidirectional: two-pass → superfluid counterflow (complex phase-conjugate unitary state)
3. Routing similarity: flat dot-product → Log-Exp deformed similarity (the pentagonator kernel)

The three substitutions

B.1 Expert gating: Log-Exp pentagonator

Replace the expert gating activation SiLU(g) with the Log-Exp deformation:

    gate' = sign(g) · [exp(β·|g|) − 4·log(1 + g²)]

    β: learnable per-expert (init 0.4, the algebra doc's constant)

The Hopfield term exp(β|g|) creates sharp attractor basins (the expert "snaps" to its specialization as the canvas denoises). The Lorentzian barrier −4log(1+g²) prevents all canvas tokens from collapsing onto the same expert (a real risk in bidirectional mode where every token sees the same context).

The pentagonator interpretation: the cluster mutation y' = y + log(1 + eˣ) on the associahedron K_5 closes the 5-fold re-bracketing. The expert gating is the deformation that lets the 512-expert routing "re-parenthesize" as the canvas denoises — each denoising step applies one cluster mutation.

B.2 GDN bidirectional: superfluid counterflow

Replace the two-pass GDN with the single-pass counterflow (the deferred fix from GemFlash38_superfluid_state_fix_DEFERRED.md):

    # per-batch, per-layer: single sweep, both modes in one update
    S_fwd = zeros(B, d, d); S_bwd = zeros(B, d, d)
    for i in 1..L_canvas:
        S_fwd = γ·S_fwd + φ(k_i) ⊗ v_i                    # forward mode, +iωt
        S_bwd = γ·S_bwd + conj(φ(k_{L-i+1})) ⊗ v_{L-i+1}  # backward mode, −iωt
        # store S_fwd[i], S_bwd[L-i+1] — both modes live in one pass

    # readout at position i sees BOTH accumulators (mutual interaction):
    out_i = softmax-norm(q_i · (S_fwd[i] + S_bwd[i])) · v

Where φ(·) is the Log-Exp deformation on the similarity score (not the features — the falsified assumption D1 from the appendix): φ(s) = sign(s) · softplus(β·|s|)

The GDN layer's gate (the "viscosity" γ) becomes a learned per-dim parameter — the superfluid model's normal-fluid fraction.

B.3 Router similarity: Log-Exp deformed

The router dot-product gate_inp · x_t → deformation on the similarity score (the pentagonator kernel from the algebra doc):

    logits = sign(s) · softplus(β · |s|)    where s = gate_inp · x_t

This replaces the flat dot-product with the Log-Exp mutation, introducing the Hopfield attractor + Lorentzian barrier into the routing itself. In bidirectional mode, this should sharpen the expert selection as the canvas converges — the "routing tightens as the denoising progresses" behavior.

Training loop

Same skeleton as Alternative A, with three differences:

Step	A (SOTA)	B (Log-Exp)
Expert gating	SiLU(g) ⊙ u	sign(g)·[exp(β·
GDN bidirectional	two-pass fwd+bwd concat	single-pass counterflow (complex conjugate modes)
Router similarity	dot-product	Log-Exp deformed similarity

Additional training considerations for B

Parameter	Value	Rationale
β init	0.4 (per-expert, learnable)	the algebra doc's constant; sharp enough to differentiate experts, soft enough for gradient flow
γ init	0.95 (per-dim, learnable)	near-unit retention; the GDN gate controls the "viscosity"
GDN bidirectional state	2 × d² per batch elt (complex conj)	single pass, not two
Extra params	β per expert + γ per GDN dim	~0.3M added (negligible)
Warmup	longer than A (β and γ need to stabilize)	500–800 steps vs 300–500

Compute (vs A)

Component	A	B	Delta
Expert gating	SiLU (trivial)	exp + log per element	+~5% expert FF
GDN	two-pass linear scan	single-pass complex state	−~40% GDN FLOPs (single vs double scan)
Router	dot-product	deformed dot-product	+~2%
Net	—	—	−~15% GDN, +~2% expert → net faster

The single-pass counterflow is faster than the two-pass bidirectional GDN (Alternnative A's approach) because it halves the sequential scan. The Log-Exp gating adds element-wise ops to the expert FF — negligible against the matmul.

Expected outcome (research hypothesis)

The Log-Exp pentagonator should:

1. Prevent the bidirectional-mode expert collapse that Alternative A is prone to (the Lorentzian barrier maintains routing diversity across the canvas)
2. Recover the fold quality cost faster (the Hopfield attractor sharpens the blended experts' specialization as the canvas denoises)
3. Produce a smoother diffusion trajectory (the pentagonator's smooth re-bracketing vs the hard causal→bidirectional transition)

These are hypotheses — the 35B-A3B benchmark is the arbiter.

Additional risk

The Log-Exp gating changes the expert function. The v3.5 expert-function distillation could be used here: first heal the experts for the bidirectional objective (Alternative A), then swap in the Log-Exp gating and fine-tune briefly. Two-phase: (1) A recovers the function, (2) B replaces the activation.

---

Comparison summary

	Alternative A (SOTA)	Alternative B (Log-Exp)
Novelty	none (standard diffusion-MoE)	high (the research formula)
Risk	low (established methods)	medium (unproven activation, unproven GDN)
Compute	baseline	~15% faster (single-pass GDN)
Training	standard pipeline	requires Log-Exp integration
Quality	expected: recovers to ~9.2–9.5 ppl	hypothesis: recovers faster + maintains routing diversity
For the paper/blog	supporting evidence	the novel contribution

Recommendation

Run A first as the control (established methods, reliable), then B as the research delta (the Log-Exp substitutions) — measuring the 35B-A3B benchmark on both. The delta between A and B is the research contribution. The v3.5 expert-function distillation (alternative path) can serve as the healing stage before either A or B, reducing the starting damage from +9.4% to a lower base.

Has in two weeks evolved to: https://github.com/ganzuul/LaserCortex/tree/main/LaserCortex/ConditionalMemory

Custom GEMM kernel exists; research grant at LUMI-G awarded for tests; attempting responsible disclosure scheme where Lean 4 proof are public while operational code is private during the experiments.