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.

Was anything tested?

No, none of the theories in these gists appear to be empirically tested yet.

They are strictly theoretical blueprints, conceptual notes, and draft research plans. The status headers inside the files explicitly label them as Status: draft (research, not yet scheduled) and mark implementation fixes as DEFERRED. Furthermore, there are no benchmark plots, loss curves, training logs, or code executions included in the gists—and the file naming convention (GemFlash38... in a workspace named llocollama) indicates these were synthesized during conceptual brainstorming sessions with an LLM.

---

What would it take to actually test this?

Turning these fluid-dynamics-inspired attention theories into working software would require a non-trivial AI research and engineering effort. Here is the concrete roadmap for what testing this concept would entail:

1. Mathematical Discretization & Formalization

• From Continuous to Discrete: Physical fluid mechanics use continuous partial differential equations (PDEs). Machine learning operates on discrete matrices. The author must first write down the exact forward and backward pass equations that map "superfluid counterflow" into standard tensor operations (Query, Key, Value matrices).
• Lie Group / Log-Exp Operations: For the mentioned "Log-Exp algebra," the transformations need to be explicitly mapped onto Lie groups so that vector operations preserve the desired geometric properties on a continuous manifold.

2. Custom Triton or CUDA Kernels

• Standard PyTorch functions (torch.nn.MultiheadAttention) and hardware accelerators are optimized for standard matrix multiplication and traditional causal/bidirectional masks.
• Standard FlashAttention won't support coupled two-way "counterflow" out of the box. To achieve acceptable GPU speed and memory efficiency, someone would have to write custom Triton or CUDA kernels to execute the two-way attention pass simultaneously on GPU SRAM.

3. Integration into a Model Architecture

• The mechanism would need to be plugged into a non-autoregressive architecture, such as a Diffusion Transformer (DiT), a Continuous Normalizing Flow (CNF), or a Masked Language Model (like BERT or a bidirectional canvas generator).
• Standard autoregressive text models (like GPT) only look backward, so a counterflow mechanism would only make sense in models where future and past context exist simultaneously on a generation canvas.

4. Compute & Experimental Benchmarking

• Baseline Setup: Train two identical model architectures on the same dataset (e.g., image generation via Flow Matching, or sequence modeling via non-autoregressive diffusion). One model uses standard multi-head attention; the other uses the proposed "superfluid counterflow" attention.
• Compute Required: Testing this effectively would require moderate-to-large GPU compute (e.g., an 8x H100/A100 node) for several days to achieve convergence on meaningful benchmark datasets.

5. Validation Metrics

To prove the theory works, the experimenters would need to measure:

• Perplexity / Fréchet Inception Distance (FID): Is the generation quality higher or more coherent?
• Convergence Speed: Does the "superfluid" dynamic allow the model to reach lower loss in fewer training steps/denoising timesteps?
• Memory Bandwidth Efficiency: Does the counterflow design reduce the memory bottleneck (the "viscous drag") during the attention computation?

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.