gamer729/cognithread-ssm

CogniThread SSM Architecture

JAX/Flax implementation of the CogniThread SSM architecture with WaveScalar thermal-analog encoding.

Architecture

1. WaveScalar — Thermal Analog Wave Representation

Each scalar is encoded as a radionics-like scalar wave with 5 parameters:

• Amplitude (A): peak value carrier
• Phase Shift (φ): initial phase offset
• Frequency (f): oscillation frequency in Hz
• Period (T = 1/f): wave period (derived)
• Fractal Dimension (D): controls harmonic overtone modulation

Effective value: v(t) = A·sin(2πft + φ)·fractal_mod(t, D)

Fractal modulation adds self-similar harmonics: 1 + Σ(1/n^D)·sin(2πnft + nφ)

No bit precision (8/16/quantized) — pure wave-scalar radionics-like analog representation.

2. CogniThreadSSM — Continuous State-Space Model

• Continuous-time ODE: dh/dt = A·h(t) + B·x(t), y(t) = C·h(t)
• ZOH discretization: Ā = exp(ΔA), B̄ = (ΔA)⁻¹(exp(ΔA)-I)·ΔB
• Parallel scan via jax.lax.associative_scan — maintains h_t in SRAM only
• Selective SSM with input-dependent Δ, B, C parameters

3. FuzzyHeuristicRetriever — Gaussian Fuzzy Membership

• μ_C(h_t, S_i) = exp(-γ‖h_t - S_i‖²) — Gaussian fuzzy membership
• Top-k retrieval: O(log|S|·d) complexity
• Learnable strictness parameter γ

4. LatentDiffusionSynthesizer — Continuous-Time Latent Diffusion

• Forward: q(z_τ|z_{τ-1}) = N(z_τ; √(1-β_τ)z_{τ-1}, β_τI)
• Reverse: p_θ(z_{τ-1}|z_τ, h_t) = N(z_{τ-1}; μ_θ, Σ_θ) conditioned on SSM state
• U-Net MLP for noise prediction
• Time complexity: O(T·d_z²)

5. AdaptiveGate — Variance-Constrained Gating

• y_final = w_f·Retrieve(S_best) + (1-w_f)·Generate(z_0)
• Var(y_final) = (1-w_f)²·σ²z₀ → 0 as w_f → 1 (anti-hallucination)
• Total time complexity: O(d²) + O(log|S|·d) + O(T·d_z²)

Model Variants

Variant	d_model	d_state	n_heuristics	d_z	T_diff	GGUF Size
Small	64	8	32	16	5	~1.4 MB
Medium	128	16	64	32	10	~9.3 MB
Large	256	32	128	64	10	~69.9 MB

GGUF Format

All models saved in GGUFv3 format with:

• Wave-scalar thermal-analog tensor encoding (5 params per scalar)
• Metadata describing encoding scheme
• Standard F32 storage for each wave parameter component

Source Files

• src_wave.py — WaveScalar and WaveScalarLinear modules
• src_ssm.py — CogniThreadSSM core
• src_fuzzy.py — FuzzyHeuristicRetriever
• src_diffusion.py — LatentDiffusionSynthesizer
• src_gate.py — AdaptiveGate
• src_model.py — CogniThreadModel (full integration)
• src_train.py — PILE streaming training with wave-scalar computing
• src_gguf.py — GGUF writer with wave-scalar encoding
• src_helpers.py — Parameter extraction utilities
• src_main.py — Build, train, and export orchestration