Losion Architecture — Tri-Jalur Router v2.0.0

Comprehensive technical reference for the Losion open-source AI framework. This document covers every component in detail, with mathematical foundations, implementation specifics, and design rationale — written for both human researchers and AI agents.

Version: v2.0.0 — Alive Gradients & Production Ready

Overview
Architecture Diagram
Gradient Flow Architecture (v1.9.0)
Jalur 1: SSM Terpadu
Jalur 2: Attention + Compression
Jalur 3: Specialized Retrieval (MoE)
Adaptive Router
Evoformer Feedback (5 Levels)
Dual Memory System
Training Pipeline (4-Phase)
Inference Pipeline
Reasoning Engine
Elastic Inference
Output Pipeline
Parameter Computation
Design Decisions & Justification
Comparison with Other Architectures

1. Overview

Losion is a generative AI architecture built on the Tri-Jalur Router paradigm — three complementary computational pathways whose contributions are dynamically weighted per-token by an adaptive router. The name Tri-Jalur (Indonesian: "Three Pathways") reflects the core insight that no single computational primitive — whether attention, state-space modeling, or retrieval — is optimal for all tokens in all contexts.

v2.0.0 — Alive Gradients & Production Ready

Version 2.0.0 fixes the last remaining dead gradient path: the AuxFreeMoE MTP loss.

AuxFreeMoE: MTP Loss Propagation

Before: MTPMoEHead inside AuxFreeMoE computed mtp_loss during forward pass and stored it in auxiliary_losses["mtp_loss"]. However, LosionForCausalLMV2.forward() never extracted this loss from the routing info and added it to the model's total loss. This meant that 32.2% of model parameters (all MTPMoEHead.pred_heads tensors) received zero gradient — they were computed but never learned.

After: LosionForCausalLMV2.forward() now iterates over all layers' routing_info["retrieval_aux"] dicts, extracts "mtp_loss" tensors with requires_grad=True, averages them, and adds them to the total loss:

# v2.0.0: Propagate AuxFreeMoE MTP loss to total loss
for layer_info in routing_info_list:
    ret_aux = layer_info.get("retrieval_aux")
    if isinstance(ret_aux, dict) and "mtp_loss" in ret_aux:
        mtp_l = ret_aux["mtp_loss"]
        if mtp_l.requires_grad:
            moe_mtp_loss += mtp_l
            n_moe_mtp += 1
if n_moe_mtp > 0:
    loss += moe_mtp_loss / n_moe_mtp

This ensures that every parameter in the model — including the MTPMoEHead prediction heads that provide complementary training signal for expert specialization — now receives training gradients.

v1.9.0 — Complete Gradient Flow & Vectorized Attention

Version 1.9.0 introduces two major architectural themes:

Complete Gradient Flow: Every component in the architecture — Evoformer feedback, Dual Memory, ThinkingToggle, and all pathway sub-layers — now maintains a fully differentiable path from output to parameters. This eliminates gradient dead-ends that existed in prior versions where in-place buffer updates and detached states could block training signal.
Vectorized Attention: Mamba-2 uses cumsum-based parallel scan (no Python loop), RWKV-7 uses parallel cumsum WKV, and Lightning Attention uses vectorized pair_mask computation. These replace sequential Python loops with batched tensor operations for 3-10× training speedup.

Why Three Pathways?

Traditional transformer architectures (e.g., GPT) apply a single uniform mechanism — self-attention — to all tokens. This is suboptimal because:

Sequential dependencies (syntax, temporal patterns) do not require O(n²) attention; they are better modeled by O(n) state-space models.
Reasoning (logical inference, multi-step comparison) genuinely benefits from full attention but is computationally expensive.
Factual knowledge is better served by explicit retrieval (MoE + engram) than by distributing facts across dense attention weights.

Losion resolves this by assigning each type of computation to a pathway optimized for it:

Pathway	Name	Optimized For	Complexity	Key Innovation
1	SSM Terpadu	Long-range sequential dependencies	O(n)	Mamba-2 SSD + Mamba-3 + RWKV-7 + Routing Mamba + Liquid SSM + PoST Decay + FG2-GDN + Structured Sparse + DeltaNet interleaved
2	Attention + Compression	Reasoning with memory efficiency	O(n·d_latent)	MLA+KDA compression (8× savings), Lightning Attention, Gated Attention, MoBA, iRoPE
3	Specialized Retrieval	Factual & domain-specific knowledge	O(n·k) sparse	AuxFreeMoE + S'MoRE + Symbolic-MoE + ∞-MoE + Engram Memory

Core Design Principles

Hardware-Agnostic: Pure PyTorch — runs on NVIDIA (CUDA) and AMD (ROCm) without code changes. torch.compile() provides graph optimization without custom kernels.
Aux-Loss-Free: Router uses bias-based routing (DeepSeek-V3 style) — no auxiliary loss needed for load balancing.
Adaptive Computation: Router adjusts compute per-token based on complexity; thinking mode activates deeper processing via sigmoid soft-blending (fully differentiable).
Memory-Efficient: MLA KV compression (8× reduction) + SSM linear recurrence + progressive KV compression + Dual Memory system.
Scalable: From 1B (prototype) to 48B+ (production) with identical architecture, varying only d_model, n_layers, and expert counts.
Inference-Scalable: MCTS, parallel thinking, and neuro-symbolic verification allow trading compute for quality at inference time.
Complete Gradient Flow: All feedback loops (Evoformer, DualMemory, ThinkingToggle) maintain differentiable paths from loss to parameters.

Agent Context: Losion is configured via LosionConfig (see losion/config.py). Sub-configs: SSMConfig, AttentionConfig, RetrievalConfig, RouterConfig, AttnResConfig, EvoformerConfig, Child3WConfig, AnchoredDecoderConfig, DualMemoryConfig, OutputConfig, JEPAConfig, DAPOConfig, RLVRConfig, PrefetchConfig, TrainingConfig, HardwareConfig, QuantizationConfig. Entry point: LosionModelV2 (see losion/models/losion_model_v2.py).

2. Architecture Diagram

Complete data flow through a single LosionLayer with all v1.9.0 components:

                            INPUT x [B, S, d_model]
                                       │
                      ┌────────────────┼────────────────┐
                      │                │                 │
                ┌─────▼─────┐   ┌─────▼─────┐   ┌──────▼──────┐
                │  SSM Norm  │   │ Attn Norm  │   │ Retr. Norm  │
                │  (RMSNorm) │   │ (RMSNorm)  │   │  (RMSNorm)  │
                └─────┬─────┘   └─────┬─────┘   └──────┬──────┘
                      │                │                 │
           ┌──────────▼──────────┐     │                 │
           │   JALUR 1: SSM      │     │                 │
           │   Terpadu Layer     │     │                 │
           │  ┌───────────────┐  │     │                 │
           │  │ Interleaving  │  │     │                 │
           │  │ Scheduler     │  │     │                 │
           │  │ (4:1:1 ratio) │  │     │                 │
           │  └───────┬───────┘  │     │                 │
           │          │          │     │                 │
           │  ┌───────▼───────┐  │     │                 │
           │  │  ┌─────────┐  │  │     │                 │
           │  │  │ Mamba-2  │  │  │     │                 │
           │  │  │   SSD    │──┤  │     │                 │
           │  │  └─────────┘  │  │     │                 │
           │  │  ┌─────────┐  │  │     │                 │
           │  │  │ Mamba-3  │  │  │     │                 │
           │  │  │  (half   │──┤  │     │                 │
           │  │  │   state) │  │  │     │                 │
           │  │  └─────────┘  │  │     │                 │
           │  │  ┌─────────┐  │  │     │                 │
           │  │  │ RWKV-7   │  │  │     │                 │
           │  │  │   WKV    │──┤  │     │                 │
           │  │  └─────────┘  │  │     │                 │
           │  │  ┌─────────┐  │  │     │                 │
           │  │  │ Routing  │  │  │     │                 │
           │  │  │ Mamba    │──┤  │     │                 │
           │  │  └─────────┘  │  │     │                 │
           │  │  ┌─────────┐  │  │     │                 │
           │  │  │ Liquid   │  │  │     │                 │
           │  │  │ SSM      │──┤  │     │                 │
           │  │  └─────────┘  │  │     │                 │
           │  │  ┌─────────┐  │  │     │                 │
           │  │  │ PoST     │  │  │     │                 │
           │  │  │ Decay    │──┤  │     │                 │
           │  │  └─────────┘  │  │     │                 │
           │  │  ┌─────────┐  │  │     │                 │
           │  │  │ Str.Sparse──┤  │     │                 │
           │  │  │ SSM      │  │  │     │                 │
           │  │  └─────────┘  │  │     │                 │
           │  │  ┌─────────┐  │  │     │                 │
           │  │  │ FG2-GDN  │  │  │     │                 │
           │  │  │(Fine-Gr. │──┤  │     │                 │
           │  │  │ Gated ΔN)│  │  │     │                 │
           │  │  └─────────┘  │  │     │                 │
           │  │  ┌─────────┐  │  │     │                 │
           │  │  │ Gated    │  │  │     │                 │
           │  │  │DeltaNet  │──┤  │     │                 │
           │  │  └─────────┘  │  │     │                 │
           │  └───────┬───────┘  │     │                 │
           │          │          │     │                 │
           │    ssm_out [B,S,D]  │     │                 │
           └──────────┬──────────┘     │                 │
                      │                │                 │
                      │    ┌───────────▼───────────┐     │
                      │    │  JALUR 2: Attention +  │     │
                      │    │     Compression        │     │
                      │    │  ┌─────────────────┐  │     │
                      │    │  │  MLA + KDA       │  │     │
                      │    │  │  (KV Compress +  │  │     │
                      │    │  │   Key-Dep Attn)  │  │     │
                      │    │  └────────┬────────┘  │     │
                      │    │           │           │     │
                      │    │  ┌────────▼────────┐  │     │
                      │    │  │ Lightning Attn  │  │     │
                      │    │  │ (vectorized     │  │     │
                      │    │  │  linear+local)  │  │     │
                      │    │  └────────┬────────┘  │     │
                      │    │           │           │     │
                      │    │  ┌────────▼────────┐  │     │
                      │    │  │ Gated Attention │  │     │
                      │    │  │ (Qwen sigmoid)  │  │     │
                      │    │  └────────┬────────┘  │     │
                      │    │           │           │     │
                      │    │  ┌────────▼────────┐  │     │
                      │    │  │  MoBA           │  │     │
                      │    │  │  (Block Attn)   │  │     │
                      │    │  └────────┬────────┘  │     │
                      │    │           │           │     │
                      │    │  ┌────────▼────────┐  │     │
                      │    │  │  Child-3W       │  │     │
                      │    │  │  (QKV MoE)      │  │     │
                      │    │  └────────┬────────┘  │     │
                      │    │           │           │     │
                      │    │  ┌────────▼────────┐  │     │
                      │    │  │  AttnRes        │  │     │
                      │    │  │  (Attn Residual)│  │     │
                      │    │  └────────┬────────┘  │     │
                      │    │           │           │     │
                      │    │  ┌────────▼────────┐  │     │
                      │    │  │ Shared Attention│  │     │
                      │    │  │ (Zamba2-style)  │  │     │
                      │    │  └────────┬────────┘  │     │
                      │    │           │           │     │
                      │    │  ┌────────▼────────┐  │     │
                      │    │  │ Cross-Jalur     │  │     │
                      │    │  │ Routing          │  │     │
                      │    │  └────────┬────────┘  │     │
                      │    │           │           │     │
                      │    │  ┌────────▼────────┐  │     │
                      │    │  │ Context Ext.    │  │     │
                      │    │  │ + iRoPE         │  │     │
                      │    │  └────────┬────────┘  │     │
                      │    │           │           │     │
                      │    │   attn_out [B,S,D]    │     │
                      │    └───────────┬───────────┘     │
                      │                │                 │
                      │                │    ┌────────────▼────────────┐
                      │                │    │  JALUR 3: Specialized   │
                      │                │    │      Retrieval          │
                      │                │    │  ┌──────────────────┐   │
                      │                │    │  │  Engram Memory   │   │
                      │                │    │  │  (Hash-based     │   │
                      │                │    │  │   Fact Store)    │   │
                      │                │    │  └────────┬─────────┘   │
                      │                │    │           │             │
                      │                │    │  ┌────────▼─────────┐   │
                      │                │    │  │  MoE Pool:       │   │
                      │                │    │  │  AuxFreeMoE      │   │
                      │                │    │  │  S'MoRE          │   │
                      │                │    │  │  Symbolic-MoE    │   │
                      │                │    │  │  ∞-MoE           │   │
                      │                │    │  │  MoHGE           │   │
                      │                │    │  │  Expert Choice   │   │
                      │                │    │  │  Gradient Routed │   │
                      │                │    │  │  Matryoshka MoE  │   │
                      │                │    │  │  Asymmetric MoE  │   │
                      │                │    │  │  Heterogeneous   │   │
                      │                │    │  └────────┬─────────┘   │
                      │                │    │           │             │
                      │                │    │  ┌────────▼─────────┐   │
                      │                │    │  │  Gated Fusion    │   │
                      │                │    │  │  (Engram + MoE)  │   │
                      │                │    │  └────────┬─────────┘   │
                      │                │    │           │             │
                      │                │    │  retrieval_out [B,S,D] │
                      │                │    └───────────┬─────────────┘
                      │                │                │
                ┌─────▼────────────────▼────────────────▼─────┐
                │              MERGE + ROUTING                  │
                │                                              │
                │  ┌─────────────────────────────────────┐    │
                │  │        Adaptive Router               │    │
                │  │  ┌──────────┐  ┌──────────────────┐  │    │
                │  │  │  Bias    │  │  Thinking        │  │    │
                │  │  │  Router  │  │  Toggle          │  │    │
                │  │  └────┬─────┘  └────┬─────────────┘  │    │
                │  │       │             │                 │    │
                │  │  ┌────┴─────┐      │                 │    │
                │  │  │Symbolic  │      │                 │    │
                │  │  │  MoE     │      │                 │    │
                │  │  └────┬─────┘      │                 │    │
                │  │       └──────┬──────┘                 │    │
                │  │              │                        │    │
                │  │     routing_weights [B, S, 3]         │    │
                │  │     [w_ssm, w_attn, w_retr]           │    │
                │  └──────────────┬────────────────────────┘    │
                │                 │                              │
                │  ┌──────────────▼────────────────────────┐    │
                │  │      Evoformer Co-evolution (L5)       │    │
                │  │   RouterExpertCoevolve adjustment      │    │
                │  └──────────────┬────────────────────────┘    │
                │                 │                              │
                │  merged = w_ssm·ssm_out +                     │
                │          w_attn·attn_out +                     │
                │          w_retr·retrieval_out                  │
                │                                              │
                │  output = x + merged  (residual)              │
                │  output = PostMergeNorm(output)               │
                └──────────────────────┬────────────────────────┘
                                       │
                            OUTPUT [B, S, d_model]
                                       │
                 ┌─────────────────────┼───────────────────────┐
                 │                     │                        │
           ┌─────▼─────┐        ┌─────▼─────┐          ┌──────▼──────┐
           │  Evoformer │        │  Dual      │          │   Output    │
           │  Feedback  │        │  Memory    │          │   Pipeline  │
           │ (5 Levels) │        │ (Working + │          │ (L-MTP +    │
           │  LayerRecy │        │  Long-term)│          │  Speculative│
           │  TokenRecy │        │            │          │  + Mirror + │
           │  DecFeedbk │        │            │          │  Anchored   │
           │  PredRecy  │        │            │          │  Diffusion) │
           │  RouterCoE │        │            │          │             │
           └───────────┘        └───────────┘          └──────┬──────┘
                                                               │
                                                    LOGITS [B, S, V]

Agent Context: The LosionLayer class in losion/models/losion_model.py and LosionModelV2 in losion/models/losion_model_v2.py implement this entire flow. Each layer instantiates SSMTerpaduLayer, AttentionKompresiLayer, RetrievalTerpaduLayer, and AdaptiveRouter. The LosionModelV2 adds Evoformer feedback, AttnRes, Dual Memory, and enhanced output pipeline.

3. Gradient Flow Architecture (v1.9.0)

Version 1.9.0 introduces complete gradient flow through all pathways, eliminating dead gradient paths that existed in previous versions.

3.1 The Gradient Flow Problem

In v0.3–v1.8, several components used in-place buffer updates or detached tensors for stability during training, which inadvertently blocked gradient flow:

Evoformer LayerRecycling: Revision signal was only applied to shallow layers; deep layers (including the final output) had no gradient path through revision parameters.
Evoformer RouterExpertCoevolve: In-place nn.Parameter updates with torch.no_grad() meant expert_state_update and state_gate parameters received no gradient.
DualMemory: LongTermMemory.consolidate() updated a buffer in-place; the output_proj and state_proj parameters could receive no gradient from the buffer.
ThinkingToggle: The depth_multiplier used a hard step function, preventing gradient flow to the complexity estimation network.

3.2 v1.9.0 Gradient Flow Fixes

Evoformer: LayerRecycling

Before: Revision applied to shallow layers only; recycled[-1] (the output used by the model) had no revision gradient.

After: Deep layers also receive a residual revision signal:

# v1.9.0: Deep layers receive revision residual
for i, h in enumerate(hidden_states):
    if i < mid:
        revised.append(h + revision * (0.1 if i < mid // 2 else 0.2))
    else:
        revised.append(h + revision * 0.05)  # NEW: deep layers too

This ensures recycled[-1] carries gradient through the revision path to all layer_recycling parameters.

Evoformer: RouterExpertCoevolve

Before: In-place self.coevolve_state.data[...] = ... with torch.no_grad() blocked all gradient.

After: update_state() returns the differentiable update tensor. The forward pass accumulates a tiny contribution to preserve the gradient path:

# v1.9.0: Differentiable path through co-evolve
total_update = torch.zeros(1, device=routing_weights.device)
for idx, output in enumerate(pathway_outputs):
    if idx < self.num_pathways:
        update = self.update_state(idx, output)  # Returns differentiable tensor
        total_update = total_update + update.sum() * 0.001
adjusted = adjusted + total_update.unsqueeze(-1) * 0  # Preserve grad, zero contribution

The in-place state update still uses detached values for stability, but the returned tensor preserves the differentiable path through expert_state_update and state_gate.

DualMemory: Direct Differentiable Path

Before: LongTermMemory.retrieve() projected the stale buffer through output_proj, but the buffer content was detached.

After: DualMemorySystem.read() establishes a direct differentiable path:

# v1.9.0: Direct differentiable path through LTM params
# Process current input x (not detached) through consolidation pipeline
x_pooled = x.mean(dim=(0, 1))
ltm_direct = self.long_term_memory.output_proj(
    self.long_term_memory.state_proj(x_pooled)
)
return x + 0.05 * memory_context + 0.01 * ltm_direct

This ensures key_proj, value_proj, query, state_proj, and output_proj all receive gradients during training.

ThinkingToggle: Sigmoid Soft-Blending

Before: Hard step function — depth_multiplier = 1.0 + complexity if complexity > threshold else 1.0 — zero gradient when below threshold.

After: Smooth sigmoid blending makes depth_multiplier fully differentiable:

# v1.9.0: Fully differentiable depth_multiplier
depth_multiplier = 1.0 + complexity * sigmoid(W_blend · x)  # Smooth blend

The sigmoid function provides non-zero gradient everywhere, allowing the complexity estimation network to be trained end-to-end.

3.3 Entropy Regularization: All Layers

In v1.9.0, entropy regularization is applied from ALL layers, not just layer 0. This prevents entropy collapse in deep layers:

$$\mathcal{L}_\text{entropy} = -\lambda_e \sum_{l=0}^{L-1} \frac{1}{L} H(\pi_l)$$

where $H(\pi_l) = -\sum_a \pi_l(a) \log \pi_l(a)$ is the entropy of the output distribution at layer $l$, and $\lambda_e$ is the entropy coefficient.

4. Jalur 1: SSM Terpadu

Jalur 1 is the sequential processing backbone of Losion. Instead of using a single SSM variant, Losion combines nine innovations in one coherent SSMTerpaduLayer with a configurable interleaving pattern.

4.1 Mamba-2 SSD (State Space Duality)

Mamba-2 SSD implements the State Space Duality principle from Gu & Dao (2024).

v1.9.0 Optimization: Uses cumsum-based parallel scan (no Python loop) for 3-10× training speedup:

# v1.9.0: Vectorized cumsum-based parallel scan
dA = torch.exp(dt.unsqueeze(-1) * A)  # [B, S, d_state]
dB = dt.unsqueeze(-1) * B              # [B, S, d_state]
# Cumulative product for state evolution
dA_cumprod = torch.cumprod(dA, dim=1)  # Parallel scan via cumsum in log space
# ... vectorized output computation

Discrete SSM formulation:

$$h_t = \bar{A} \cdot h_{t-1} + \bar{B} \cdot x_t$$

$$y_t = \bar{C} \cdot h_t$$

Default chunk size: 256 tokens (balances parallelism and memory usage).

Agent Context: Implemented in Mamba2SSD class (lossion/core/ssm/mamba2.py). Key params: d_model, d_state (default 64), d_conv (default 4), expand (default 2), chunk_size (default 256).

4.2 Mamba-3 SSD — Inference-First Design

Mamba-3 (arXiv:2603.15569) is an evolution of Mamba-2 with three key improvements from an inference-first perspective:

Reduced state dimension (d_state=32 vs Mamba-2's 64): Half the state size but with better utilization through optimized S4D initialization and dual token shift.
Dual Token Shift: Two independent shift patterns (inspired by RWKV) — forward shift and backward shift — combined via learnable mixing coefficient:
```
alpha = torch.sigmoid(self.mix_alpha)
mixed = alpha * shift_fwd_proj(x_fwd) + (1 - alpha) * shift_bwd_proj(x_bwd)
return x + mixed
```

Inference-first dt discretization: Clamped exponential and softplus-stabilized input scaling prevent numerical instability on long sequences:

dt_clamped = dt.clamp(min=0.0, max=dt_clamp_max)
dA = torch.exp((dt_clamped.unsqueeze(-1) * A).clamp(max=0.0))  # Always ≤ 1
dt_stabilized = F.softplus(dt_clamped)
dB = dt_stabilized.unsqueeze(-1) * B

Agent Context: Mamba3SSD class (lossion/core/ssm/mamba3.py). Default d_state=32. Includes DualTokenShift module.

4.3 RWKV-7 WKV (Weighted Key-Value)

RWKV-7 implements WKV recurrence — an evolution of attention that replaces softmax scoring with an exponentially-weighted moving average.

v1.9.0 Optimization: Uses parallel cumsum WKV instead of sequential Python loop:

# v1.9.0: Vectorized WKV via parallel cumsum
decay = torch.exp(w)  # [B, S, H] — data-dependent decay
# Cumulative product of decay for parallel scan
decay_cumprod = torch.cumprod(decay, dim=1)
# ... vectorized numerator/denominator computation

Key advantages:

O(1) inference: Only stores (wkv_state, sum_state)
Unbounded context: State accumulates without sequence length limit
Explicit forgetting: Decay factor $a_t = \exp(w_t)$ allows selective "forgetting"

Agent Context: RWKV7WKV class (lossion/core/ssm/rwkv7.py). The u parameter (position bonus) is learned per-head.

4.4 Routing Mamba (RoM) — MoE over SSM Projections

Routing Mamba from Microsoft Research (NeurIPS 2025) scales SSM parameters using sparse mixtures of linear projection experts. Instead of a single set of B, C, dt projections, Routing Mamba maintains multiple expert projection sets and routes each token to a sparse subset:

$$B_\text{eff} = \sum_{k} w_k \cdot B_{\text{expert}_k}(x), \quad C_\text{eff} = \sum_{k} w_k \cdot C_{\text{expert}_k}(x), \quad dt_\text{eff} = \sum_{k} w_k \cdot dt_{\text{expert}_k}(x)$$

Load balancing: DeepSeek-V3 aux-loss-free approach — EMA tracking of expert load with periodic bias updates (non-gradient).

Shared parameters: A_log matrix and D skip connection are shared across all experts.

Agent Context: RoutingMamba class (lossion/core/ssm/routing_mamba.py). Config: routing_mamba_num_experts (default 4), routing_mamba_active_experts (default 2).

4.5 Liquid SSM — Adaptive Compute Depth

Liquid SSM extends Mamba-2 with input-adaptive time constants controlled by a ComplexityGate:

ComplexityGate estimates per-token, per-head complexity and maps it to one of three depth levels:

Depth 1 (fast): Single SSD pass — minimal compute for easy tokens
Depth 2 (standard): SSD + one additional sub-layer
Depth 3 (deep): Full interleaving through all sub-layers

Liquid time-constant rule:

$$dt_\text{eff} = dt_\text{base} \cdot (1 + s \cdot (2 \cdot \sigma(c_\text{inner}) - 1))$$

where $s$ is a learnable scale (initialized near 0) and $c_\text{inner}$ is the complexity projected to d_inner dimensions.

Training: Soft blending of all depth outputs with depth probabilities (for gradient flow). Inference: Hard early exit for easy tokens.

Agent Context: LiquidSSMTerpaduLayer class (lossion/core/ssm/liquid_ssm.py). Includes ComplexityGate and LiquidSSD. Depth entropy loss: depth_entropy_weight (default 0.01).

4.6 PoST Decay — Position-Dependent Decay Spectra

PoST (Position-Dependent Decay Spectra) replaces the single learnable decay parameter per head with a spectrum of decay rates that vary by position:

$$h_t = \sum_m \text{mix}_t(m) \cdot \gamma(m) \cdot \text{dA}_t \cdot h_{t-1,m} + \text{dB}_t \cdot x_t / M$$

where:

$\gamma(m) \in (0,1)$: Learnable decay rate per mode (sigmoid of log_gamma)
$\text{mix}_t(m)$: Position-dependent mixing weights (softmax of position embedding + MLP)

This allows the SSM to retain information long-term (slow decay modes) while focusing on local context (fast decay modes) — position-dependent.

Agent Context: PoSTDecaySSM class (lossion/core/ssm/post_decay.py). Includes DecaySpectrum submodule. Default n_decay_modes=4.

4.7 Structured Sparse Transition SSM

Based on NeurIPS 2025 (poster 118046), this replaces the diagonal transition matrix with structured off-diagonal elements enabling FSA (Finite State Automata) state tracking.

Key insight: Diagonal SSMs can only track O(log S) FSA states. Structured sparse transitions can track O(S) states — provably optimal.

Three sparsity patterns:

Block-diagonal: Dense B×B blocks, diagonal across blocks. Complexity O(S × B).
Banded: Diagonal ± band_width. Complexity O(S × bandwidth).
Butterfly: Recursive butterfly factorization. Complexity O(S × log₂(B)).

The transition is applied via first-order matrix exponential approximation:

$$\exp(dt \cdot A) @ h \approx h + dt \cdot (A @ h)$$

Agent Context: StructuredSparseSSM class (lossion/core/ssm/structured_sparse.py). Config: n_groups (default 4), transition_type.

4.8 FG2-GDN — Fine-Grained Gated DeltaNet

FG2-GDN enhances the standard GatedDeltaNet with per-head, per-position gating:

$$\beta_t[h, s] = \text{gate_fn}(h_s \cdot W_\beta + b_\beta[h]) / \text{temperature}[h]$$ $$\alpha_t[h, s] = \text{gate_fn}(h_s \cdot W_\alpha + b_\alpha[h] + \text{offset}[h]) / \text{temperature}[h]$$

where $h$ is the head index and $s$ is the position index. This provides far more granular retention control than standard DeltaNet (which uses a single gate value per head).

Learnable temperature per head controls selectivity: low temperature → sharp (selective), high temperature → smooth (uniform).

Agent Context: FG2GDN class (lossion/core/ssm/fg2_gdn.py). Includes FineGrainedGate. Gate types: "sigmoid" or "softmax". Position bias optional.

4.9 Gated DeltaNet (Original)

The standard Gated DeltaNet implements in-context learning via the delta rule:

$$V_t = \alpha_t \cdot V_{t-1} + \beta_t \cdot (k_t^\top \cdot v_t)$$

Unlike standard attention that only adds new information, DeltaNet can correct previously stored associations.

Agent Context: GatedDeltaNet class (lossion/core/ssm/delta_net.py). Alpha initialized near 1.0 via alpha_offset.

4.10 Interleaving Pattern

The SSM sub-layers are interleaved with a default 4:1:1 ratio:

4 blocks Mamba-2 SSD: Primary parallel sequential processing
1 block RWKV-7 WKV: Dynamic state evolution with explicit forgetting
1 block Gated DeltaNet: In-context learning and state correction

Additional sub-layers (Mamba-3, Routing Mamba, Liquid SSM, PoST Decay, Structured Sparse, FG2-GDN) can be enabled via config and participate in the interleaving schedule.

Dynamic routing mode: When routing weights are provided, all sub-layers process the input simultaneously and outputs are blended:

$$\text{blended} = w_\text{ssd} \cdot y_\text{ssd} + w_\text{wkv} \cdot y_\text{wkv} + w_\text{delta} \cdot y_\text{delta} + \cdots$$

Agent Context: InterleavingScheduler in lossion/core/ssm/ssm_layer.py. State tracked via SSMState.

5. Jalur 2: Attention + Compression

Jalur 2 is the reasoning engine of Losion, combining multiple attention mechanisms with MLA compression for memory efficiency.

5.1 MLA + KDA (Multi-head Latent Attention + Key-Dependent Attention)

MLA, adapted from DeepSeek-V3, compresses KV-cache representations into a lower-dimensional latent space (8× compression). KDA adds key-dependent attention bias for improved routing information.

KV Compression math:

$$c_{kv} = W_\text{kv_compress}(x) \in \mathbb{R}^{d_\text{latent}}$$

$$k = W_\text{k_up}(c_{kv}), \quad v = W_\text{v_up}(c_{kv})$$

Memory savings: For Losion-7B with n_heads=16, d_kv=128, mla_latent_dim=512:

Full KV: $2 \times 16 \times 128 = 4{,}096$ dimensions → MLA latent: $512$ → 8× compression

Agent Context: MLA class (lossion/core/attention/mla.py). KDA class (lossion/core/attention/kda_mla.py).

5.2 Lightning Attention — Vectorized Linear + Local

Lightning Attention combines linear attention (for global context) with local windowed attention (for positional precision), both computed in a vectorized manner.

v1.9.0: Uses vectorized pair_mask computation instead of sequential masking:

# Vectorized pair_mask for causal + local window
pair_mask = causal_mask & distance_mask  # Both pre-computed as [S, S] bool tensors

Architecture:

Global component: Linear attention with feature map (elu, relu, or cos)
Local component: Standard softmax attention within sliding window
Chunk-based parallel training: Processes sequence in chunks for memory efficiency

Agent Context: LightningAttention class (lossion/core/attention/lightning_attention.py). Config: lightning_window_size (default 2048), lightning_chunk_size (default 4096).

5.3 Gated Attention (Qwen-style)

Adapted from Qwen (NeurIPS 2025 Best Paper), Gated Attention applies a sigmoid gate to the attention output:

$$\text{output} = \sigma(W_\text{gate} \cdot x) \odot \text{attention}(Q, K, V)$$

The sigmoid gate provides smooth, differentiable control over how much attention information flows to the output. Unlike ReLU or hard gates, sigmoid ensures non-zero gradient everywhere.

Agent Context: GatedAttention class (lossion/core/attention/gated_attention.py). Enabled via use_gated_attention config.

5.4 MoBA — Mixture of Block Attention

MoBA (Moonshot AI, NeurIPS 2025) partitions the sequence into blocks and routes each query to a top-K subset of blocks:

$$\text{MoBA}(q_t) = \sum_{b \in \text{TopK}_K(S(q_t, \text{blocks}))} \text{Attn}(q_t, K_b, V_b)$$

This reduces attention complexity from O(n²) to O(n × K × block_size) while preserving the ability to attend to relevant distant context.

Agent Context: MoBA class (lossion/core/attention/moba.py). Config: moba_block_size (default 512), moba_top_k_blocks (default 4).

5.5 Child-3W — QKV-Level MoE Routing

Child-3W applies MoE routing at the QKV level: multiple child attention parameter sets with routing between them. Each child has its own Q, K, V projections, and a router selects which children to activate per token.

$$\text{output} = \sum_{c \in \text{TopK}} g_c \cdot \text{Attn}(Q_c, K_c, V_c)$$

where $g_c$ is the routing weight for child $c$.

Agent Context: Child3WConfig in lossion/config.py. Config: num_children (default 4), top_k_children (default 2).

5.6 AttnRes — Attention-Based Residuals

AttnRes (MoonshotAI 2026, v0.9) replaces fixed-weight residual connections with learned attention-based aggregation across layers. Three modes:

Full: All-layer attention residual
Block: Group layers into blocks; within-block attention residual
Hybrid: Full across blocks, block within

Also supports token compression via linear, gated, or SSM-based compression.

Agent Context: AttnResConfig in lossion/config.py. Modes: "full", "block", "hybrid".

5.7 Shared Attention (Zamba2-style)

Shared Attention, adapted from Zamba2, shares attention parameters across groups of layers, with a small ratio of unique parameters per layer. This reduces parameter count while maintaining expressiveness.

Agent Context: Config: shared_n_groups (default 1), shared_pattern ("all_shared" or "interleaved"), shared_unique_ratio (default 0.25).

5.8 Cross-Jalur Routing

Cross-Jalur Routing allows information flow between the three pathways within the attention layer. A learnable blend parameter controls how much cross-pathway information is incorporated:

$$\text{attn_enhanced} = (1 - \alpha) \cdot \text{attn_out} + \alpha \cdot \text{cross_jalur}(\text{ssm_out}, \text{retrieval_out})$$

Agent Context: CrossJalurRouting class (lossion/core/retrieval/cross_jalur_routing.py). Config: cross_jalur_blend_alpha (default 0.3), cross_jalur_graph_top_k (default 8).

5.9 Context Extension + iRoPE

Context Extension dynamically adjusts the effective context window based on input complexity.

iRoPE (Interleaved RoPE) alternates between RoPE and NoPE layers with a 3:1 ratio:

Layer:  0    1    2    3    4    5    6    7    8    ...
RoPE:   ✓    ✓    ✓    ✗    ✓    ✓    ✓    ✗    ✓    ...

RoPE provides explicit positional information; NoPE allows learned attention biases for long contexts without positional extrapolation issues.

Agent Context: InterleavedRoPE class (lossion/core/attention/irope.py). ContextExtension class (lossion/core/attention/context_extension.py).

6. Jalur 3: Specialized Retrieval (MoE)

Jalur 3 handles factual knowledge and domain-specific knowledge through a layered architecture combining multiple MoE variants with Engram Memory.

6.1 AuxFreeMoE — DeepSeek-V3 Style

AuxFreeMoE eliminates the quality-degrading auxiliary loss, replacing it with bias-based load balancing (DeepSeek-V3 style):

Mechanism:

Standard routing: logits = gate_proj(x) + bias
Top-K expert selection with renormalization
Bias updated via EMA running statistics (non-gradient): overloaded experts → negative bias, underloaded → positive bias
No auxiliary loss returned — only monitoring metrics

MTP Training Signal: AuxFreeMoE includes Multi-Token Prediction heads that predict future tokens (t+1, t+2, ..., t+n), providing a complementary training signal for expert specialization without quality degradation:

$$\mathcal{L}_\text{MTP} = \sum_{k=1}^{n} \lambda^{k-1} \cdot \text{CE}(\text{pred}_k, \text{target}_k)$$

where $\lambda = 0.5$ is the geometric decay factor.

Agent Context: AuxFreeMoE class (lossion/core/retrieval/aux_free_moe.py). Includes AuxFreeMoERouter and MTPMoEHead. Config: bias_update_rate (default 0.01).

6.2 S'MoRE — Sub-tree MoE with Residual Experts

S'MoRE (Meta, NeurIPS 2025) organizes experts in a sub-tree structure with shared residual connections:

smore_num_sub_trees (default 4): Number of shared sub-trees
smore_sub_tree_depth (default 2): Depth of each sub-tree

This allows experts to share base knowledge through the sub-tree structure while specializing at the leaves.

Agent Context: SMoRE class (lossion/core/retrieval/smore.py).

6.3 Symbolic-MoE — Skill-Based Discrete Routing

Symbolic-MoE routes tokens based on discrete skill labels rather than continuous routing weights. This provides interpretable, deterministic routing for known task types:

Each expert is associated with a symbolic skill (e.g., "math", "code", "translation")
Routing is determined by the task label, not learned affinity
Complementary to continuous routing — used for known, well-defined tasks

Agent Context: SymbolicMoE class (lossion/core/retrieval/symbolic_moe.py). Also used in the Adaptive Router for skill-based pathway selection.

6.4 ∞-MoE — Continuous Expert Space

∞-MoE replaces the discrete expert pool with a continuous expert space using a codebook + hypernetwork:

Codebook: infinite_moe_codebook_size (default 256) codes of dimension infinite_moe_code_dim (default 32)
Hypernetwork: Generates expert weights from codes via infinite_moe_hypernet_hidden (default 256) dimensional hidden layer
Low-rank residual: Optional low-rank residual for efficient expert generation

Agent Context: InfiniteMoE class (lossion/core/retrieval/infinite_moe.py). Config: use_infinite_moe, infinite_moe_code_dim, infinite_moe_hypernet_hidden.

6.5 MoHGE — Mixture of Heterogeneous Group Experts

MoHGE groups experts into heterogeneous groups with different architectures or capacities, allowing the model to balance computational efficiency with expressiveness.

Agent Context: MoHGE class (lossion/core/retrieval/mohge.py).

6.6 Other MoE Variants

Expert Choice (Google Research): Experts choose tokens — guaranteed load balance
Gradient Routed MoE: Routing weights are conditioned on gradient information
Matryoshka MoE: Variable-depth expert computation (elastic inference)
Asymmetric MoE: MoE layers placed only at specific layer indices
Heterogeneous MoE: Experts with varying dimensions

6.7 Engram Memory

Engram Memory is a hash-based fact store that stores factual knowledge explicitly:

$$\text{subject_string} \xrightarrow{\text{hash}} \text{bucket_index} \xrightarrow{\text{embedding_lookup}} \text{retrieval}$$

Advantages: O(1) retrieval, explicit storage, updateable without retraining.

Agent Context: EngramMemory class (lossion/core/retrieval/engram.py). Config: num_buckets (default 1,000,000), engram_dim (default 256).

6.8 Gated Fusion

Engram and MoE outputs are combined via gated fusion:

$$w = \text{softmax}(W_\text{fusion}([\text{engram_out}; \text{moe_out}]))$$ $$\text{fused} = w_0 \cdot \text{engram_out} + w_1 \cdot \text{moe_out}$$

Three fusion modes: "gated" (default), "additive", "concat".

7. Adaptive Router

The Adaptive Router combines three components: BiasRouter (computational allocation), ThinkingToggle (complexity detection), and Symbolic-MoE (skill-based routing).

7.1 BiasRouter — Aux-Loss-Free Routing

$$\text{logits} = W_\text{router} \cdot x + b_\text{bias}$$ $$\text{weights} = \text{softmax}(\text{logits})$$

The bias $b_\text{bias}$ is updated directly by the main loss gradient — no auxiliary loss term needed (DeepSeek-V3 approach).

Agent Context: BiasRouter class (lossion/core/router/bias_router.py). Config: num_pathways (default 3), top_k_pathways (default 2).

7.2 ThinkingToggle — Complexity Detection

v1.9.0: Uses sigmoid soft-blending for depth_multiplier (fully differentiable):

$$\text{depth_multiplier} = 1.0 + \text{complexity} \cdot \sigma(W_\text{blend} \cdot x)$$

This replaces the hard step function, providing non-zero gradient everywhere.

Effects of Thinking Mode:

Aspect	Non-Thinking	Thinking
Routing weights	Jalur 1 dominant	Jalur 2+3 activated
Interleaving ratio	5:1 (local:global)	2:1 (local:global)
Pairformer	Inactive	Active
Depth multiplier	1.0	1.0–2.0 (smooth blend)
Gradient flow	Via BiasRouter	Via BiasRouter + sigmoid

Agent Context: ThinkingToggle class (lossion/core/router/thinking_toggle.py). Output: ThinkingAssessment with mode, complexity_score, dominant_task, confidence, depth_multiplier.

7.3 Symbolic-MoE Integration

Symbolic-MoE provides skill-based discrete routing as a third signal in the router. For known task types (math, code, translation), it provides deterministic, interpretable routing that complements the learned BiasRouter and ThinkingToggle.

7.4 Thinking-Weight Adjustment

After BiasRouter, ThinkingToggle, and Symbolic-MoE produce their outputs, routing weights are adjusted via a learned network with residual connection:

adjuster_input = cat([routing_weights, complexity_score.unsqueeze(-1)], dim=-1)
adjustment = thinking_adjuster(adjuster_input)  # [B, S, 3]
adjusted = routing_weights + 0.1 * adjustment   # Residual with small scale
adjusted = softmax(adjusted, dim=-1)

Agent Context: Full routing in AdaptiveRouter class (lossion/core/router/router.py). Pathway priors: [0.4, 0.3, 0.3].

8. Evoformer Feedback (5 Levels)

Adapted from AlphaFold2's Evoformer (Nobel Prize 2024), generalized as a universal architectural principle for LLMs. Five levels of bidirectional feedback:

Level 1 — Inter-Layer Recycling

Deep layers revise shallow layer representations via cross-attention:

$$\text{revision} = \text{gate} \cdot \text{Attn}(Q=\text{shallow}, K=\text{deep}, V=\text{deep})$$

v1.9.0: Deep layers also receive revision residual (0.05 scale) so that recycled[-1] carries gradient through the revision path.

Level 2 — Bidirectional Token Update

Later tokens revise earlier token representations through bidirectional attention (applied after initial causal pass):

$$\text{revised}_t = x_t + \text{gate} \cdot \text{Attn}_\text{backward}(x_t, \text{all_tokens})$$

This is NOT BERT-style — it's iterative revision AFTER the initial forward pass, preserving autoregressive reasoning.

Level 3 — Decoder ↔ Predict Feedback

Bidirectional feedback between the decoder and prediction modules:

$$\text{updated} = \text{RMSNorm}(h + \text{gate} \cdot W_\text{feedback}(\text{decoder_out} - h))$$

Level 4 — Prediction → Context Recycling

The most revolutionary level: predicted token N can revise representations of tokens 1 through N-1:

$$\text{revised} = h + \text{gate} \cdot \text{Attn}(Q=h, K=\text{prediction}, V=\text{prediction})$$

Level 5 — Router ↔ Expert Co-Evolution

Router and experts co-evolve during training. The co-evolution state captures the "negotiation" between router choices and expert specialization:

$$\text{adjustment} = \text{Tanh}(W_\text{adj}(\text{mean}(\text{coevolve_state}))) \cdot 0.1$$

v1.9.0: update_state() returns differentiable tensor preserving gradient flow through expert_state_update and state_gate parameters, while the in-place buffer update uses detached values for stability.

Agent Context: EvoformerManager class (lossion/core/feedback/evoformer.py). Config: EvoformerConfig with n_recycling_steps (default 3), 5 toggle flags for each level.

9. Dual Memory System

Two-level memory system inspired by human memory:

Working Memory

Direct access to recent token/layer outputs (ring buffer)
High detail, limited capacity (working_memory_size, default 512)
Entries are detached for persistence across forward passes

Long-Term Memory

Compressed, persistent hidden state from consolidation
Selective, persistent, compressed (long_term_memory_dim, default 256)
Three consolidation methods: "attention" (default), "gated", "mean"

v1.9.0: DualMemorySystem.read() establishes a direct differentiable path through LTM parameters:

# Direct differentiable path: x_pooled → state_proj → output_proj
ltm_direct = self.long_term_memory.output_proj(
    self.long_term_memory.state_proj(x_pooled)
)
return x + 0.05 * memory_context + 0.01 * ltm_direct

This ensures key_proj, value_proj, query, state_proj, and output_proj all receive gradients during training.

Agent Context: DualMemorySystem class (lossion/core/memory/dual_memory.py). Includes WorkingMemory (ring buffer) and LongTermMemory (attention-gated consolidation). Config: DualMemoryConfig.

10. Training Pipeline (4-Phase)

Losion uses a 4-phase training pipeline, progressing from individual component training to advanced RL:

Phase 1: Individual Component Training

Each pathway is trained independently:

Jalur 1 (SSM): Standard language modeling with SSM-only forward pass
Jalur 2 (Attention): Standard language modeling with attention-only forward pass
Jalur 3 (MoE): Standard language modeling with MoE-only forward pass

Phase 2: Joint Training

All three pathways are trained jointly with the adaptive router:

Router learns to allocate computation across pathways
BiasRouter biases are updated via gradient (no aux loss)
Entropy regularization from ALL layers (v1.9.0)
LLM-JEPA: Predicts future latent states instead of next tokens

LLM-JEPA (v0.6):

$$\mathcal{L}_\text{JEPA} = w_\text{pred} \cdot \mathcal{L}_\text{VICReg}(\text{predictor}(h_t), h_{t+k})$$

where the predictor forecasts $k$ steps ahead in latent space, and VICReg loss ensures variance, invariance, and covariance regularization.

Agent Context: JEPAConfig in lossion/config.py. prediction_horizon (default 4), loss_type ("vicreg", "cosine", "mse"), teacher_ema_decay (default 0.996).

Phase 3: RL Fine-Tuning

DAPO or GRPO (auto-selected based on config):

DAPO (Decoupled Clip & Dynamic Sampling Policy Optimization, v0.8):

Four key improvements over GRPO:

Decoupled Clip: Asymmetric clip ratios — clip_ratio_low=0.2 (looser lower bound) and clip_ratio_high=0.28 (tighter upper bound to prevent reward hacking)
Dynamic Sampling: Filter prompts where all responses have the same reward (zero learning signal) — ~15-20% efficiency gain
Token-Level Loss: Per-token policy gradient for finer credit assignment
Overlong Filtering: Penalty for responses exceeding max_response_length

RLVR (Reinforcement Learning with Verifiable Rewards, v0.8):

Uses objective, programmable reward functions instead of learned reward models:

Math verification with configurable tolerance
Code execution verification with timeout
Format checking
Curriculum difficulty scheduling: "easy" → "medium" → "hard"

Agent Context: DAPOConfig in lossion/config.py. DAPOTrainer in lossion/training/dapo.py. RLVRConfig in lossion/config.py. RLVRTrainer in lossion/training/rlvr.py.

Phase 4: Advanced Training

Evolutionary search for architecture optimization
Active learning for data efficiency
Distillation from larger models

Training Auto-Selection

The LosionRecipe and LosionOrchestrator automatically select the appropriate RL method:

if config.dapo.enabled:
    rl_trainer = DAPOTrainer(config.dapo, model, reward_fn)
elif config.rlvr.enabled:
    rl_trainer = RLVRTrainer(config.rlvr, model, reward_fn)
else:
    rl_trainer = GRPOTrainer(config.grpo, model, reward_fn)

Agent Context: LosionOrchestrator (lossion/training/losion_orchestrator.py) manages the 4-phase pipeline. LosionRecipe (lossion/training/losion_recipe.py) provides pre-configured training recipes.

11. Inference Pipeline

11.1 Expert Prefetching — Speculating Experts

"Speculating Experts" (arXiv:2603.19289) predicts which MoE experts will be needed in subsequent layers and prefetches them, overlapping expert loading latency with ongoing computation.

Architecture:

LightweightPredictor: 2-layer MLP per layer (< 1% of single expert params) maps layer-L hidden states to layer-(L+1) expert predictions
Finite MoE mode: Predicts discrete expert indices via top-k or temperature sampling
∞-MoE mode: Predicts continuous expert codes; nearby codes in the continuous space are prefetched
Adaptive temperature: Dynamically adjusts prediction temperature based on recent accuracy (exploit when accurate, explore when not)
Accuracy tracking: Rolling precision/recall/hit-rate/coverage metrics per layer

Prefetch pipeline:

Receive hidden_states from layer L
Feed into predictor[L] → predicted expert set for L+1
Issue async prefetch (overlaps with expert compute at layer L)
At layer L+1: check if needed experts are already loaded (hit → zero latency, miss → fallback)

Agent Context: ExpertPrefetcher class (lossion/inference/expert_prefetch.py). Config: PrefetchConfig with predictor_hidden_dim (default 128), prefetch_budget (default 4), adaptive_temperature.

11.2 QuantSpec — Quantization-Aware Speculative Decoding

QuantSpec combines quantization with speculative decoding for fast, memory-efficient inference.

Agent Context: QuantSpec class (lossion/inference/quantspec.py).

11.3 Paged KV Cache with INT4 Quantization

The paged KV cache manages memory efficiently with INT4 quantization:

Paged allocation: KV cache is organized in fixed-size pages, avoiding contiguous memory requirements
INT4 quantization: KV vectors are quantized to 4-bit, reducing memory by 8× compared to FP32
MLA compatibility: Compressed latent vectors from MLA are further quantized

Agent Context: kv_cache.py (lossion/inference/kv_cache.py). KVCache class with paged allocation.

11.4 Matryoshka Elastic Inference

Matryoshka Nested Transformer enables one weight set to produce multiple valid submodels of different sizes. At inference time, a size_selector network predicts the appropriate granularity factor per token:

Simple tokens → smaller submodel (faster)
Complex tokens → full model (higher quality)
Default granularity factors: [0.25, 0.5, 0.75, 1.0]

Mix'n'Match: Different layers can use different granularity factors — e.g., early layers small (0.25), late layers full (1.0).

Agent Context: MatryoshkaFFN class (lossion/core/elastic/matryoshka.py). Also MatryoshkaMoE for elastic expert computation.

12. Reasoning Engine

Losion integrates three reasoning techniques for inference-time compute scaling:

12.1 MCTS (Monte Carlo Tree Search)

AlphaZero-inspired tree search with UCB selection:

$$\text{UCB}(s, a) = Q(s, a) + c_\text{puct} \cdot P(s, a) \cdot \frac{\sqrt{N(s)}}{1 + N(s, a)}$$

Adaptive compute budget: budget = base + (max - base) · complexity²

Agent Context: MCTSReasoner class (lossion/core/reasoning/mcts.py). Config: num_simulations (default 64), c_puct (default 1.5).

12.2 Parallel Thinking — Gemini Deep Think Style

Explores multiple reasoning paths simultaneously:

$$\text{score}_i = 0.5 \cdot \text{value}_i + 0.35 \cdot \text{consistency}_i + 0.15 \cdot \text{novelty}_i$$

Selection strategies: BEST_OF_N, MAJORITY_VOTE, WEIGHTED_MERGE, TOURNAMENT.

Agent Context: ParallelThinker class (lossion/core/reasoning/parallel_thinking.py).

12.3 Neuro-Symbolic Verification — AlphaProof Style

Combines neural generation with symbolic verification for formally correct reasoning:

Symbolic Rule Engine: Learned rule embeddings with verification networks
Error Localization: Per-token error probability
Feedback Generation: Correction signal for failed verification
Iterative Revision Loop: Up to max_revision_iterations (default 3)

Agent Context: NeuroSymbolicVerifier class (lossion/core/reasoning/neuro_symbolic.py). Output: VerificationResult with status (VERIFIED/FAILED/PARTIAL/UNSURE/NEEDS_REVISION).

13. Elastic Inference

Matryoshka / MatFormer — One Weight Set, Multiple Submodels

Nested FFN structure with Matryoshka loss:

$$\mathcal{L}_\text{matryoshka} = \frac{w_\text{mat}}{|\mathcal{F}|} \sum_{f \in \mathcal{F}} \mathcal{L}(\text{submodel}_f, \text{target}_f)$$

where $\mathcal{F}$ = {0.25, 0.5, 0.75, 1.0}.

# Matryoshka forward (from losion/core/elastic/matryoshka.py)
if factor >= 1.0:
    gate = F.silu(self.gate_proj(x))
    up = self.up_proj(x)
    output = self.down_proj(gate * up)
else:
    d_ff_active = int(factor * self.d_ff)
    gate_w = self.gate_proj.weight[:d_ff_active, :]
    up_w = self.up_proj.weight[:d_ff_active, :]
    down_w = self.down_proj.weight[:, :d_ff_active]
    gate = F.silu(F.linear(x, gate_w))
    up = F.linear(x, up_w)
    output = F.linear(gate * up, down_w)

Agent Context: MatryoshkaFFN class (lossion/core/elastic/matryoshka.py). Also AttnLoRA for elastic attention adaptation.

14. Output Pipeline

14.1 L-MTP — Leap Multi-Token Prediction

L-MTP (v0.8, NeurIPS 2025) extends standard MTP with leap scheduling — predicting tokens at non-adjacent positions:

$$\text{L-MTP: predict tokens at positions } t + \Delta_1, t + \Delta_2, \ldots$$

where $\Delta_k$ follows a "geometric", "arithmetic", or "adjacent" schedule:

Geometric (default): $\Delta_k = 2^k$ — leap 1, 2, 4, 8 tokens ahead
Arithmetic: $\Delta_k = k \cdot \text{step}$ — uniform spacing
Adjacent: $\Delta_k = k$ — standard MTP

Agent Context: LeapMTP class (lossion/core/output/leap_mtp.py). Config: leap_mtp_schedule, leap_mtp_num_leaps (default 4), leap_mtp_max_leap (default 8).

14.2 Speculative Decoder

Standard speculative decoding with draft model:

# Draft model proposes tokens → Target model verifies
draft_tokens = draft_model.generate(x, num_tokens=draft_len)
probabilities = target_model(draft_tokens)
acceptance = check_acceptance(draft_tokens, probabilities)

Agent Context: SpeculativeDecoder class (lossion/core/output/speculative_decoder.py).

14.3 Mirror Speculative

Mirror speculative decoding uses the same model as both draft and target, with different routing strategies:

Draft: SSM-dominant routing (fast, linear time)
Target: Full Tri-Jalur routing (accurate)

Agent Context: MirrorSpeculative class (lossion/core/output/mirror_speculative.py).

14.4 Anchored Diffusion Decoder

Continuous vector prediction + lightweight anchored diffusion refinement:

Predict continuous vector (not softmax → token ID)
Refine via 2-3 step anchored diffusion process
Disambiguate via attention-based disambiguation heads

Uses Evoformer feedback loop (Level 3) for iterative refinement.

Agent Context: AnchoredDiffusionDecoder class (lossion/core/output/anchored_decoder.py). Config: AnchoredDecoderConfig with n_refine_steps (default 3), disambiguation_heads (default 8).

15. Parameter Computation

For the standard Losion-7B configuration:

Component	Parameters
Embedding	`vocab_size × d_model` = 32,000 × 4,096 = 131M
SSM Terpadu (per layer)	~35M × `n_layers` = 35M × 32 = 1,120M
Attention + MLA (per layer)	~25M × `n_layers` = 25M × 32 = 800M
MoE (per layer, 64 experts)	~50M × `n_layers/2` = 50M × 16 = 800M
Router (per layer)	~0.5M × 32 = 16M
Output Pipeline	~135M
Evoformer + DualMemory	~50M
Total (7B)	~3,052M (active per token)

Due to MoE sparsity (93.75%), only ~3B parameters are active per token despite total model size of ~7B.

16. Design Decisions & Justification

Why Tri-Jalur Instead of Uniform Attention?

Criterion	Uniform Attention	Tri-Jalur
Sequential tokens	O(n²) waste	O(n) SSM — optimal
Reasoning tokens	Full attention needed	Attention pathway — optimal
Factual recall	Distributed in weights	MoE + Engram — explicit
Compute per token	Constant (expensive)	Adaptive (cheap or expensive)
Memory	Full KV cache	MLA + SSM (8× savings)

Why Bias-Based Routing Instead of Auxiliary Loss?

Auxiliary loss (Switch Transformer style) adds a separate loss term: $\mathcal{L} = \mathcal{L}\text{main} + \alpha \cdot \mathcal{L}\text{aux}$. This:

Requires hyperparameter tuning ($\alpha$)
Interferes with main loss optimization
Doesn't always produce optimal load balancing

Bias-based routing (DeepSeek-V3 style) updates bias directly via gradient or EMA — no separate loss term, no quality degradation.

Why Evoformer Instead of Standard Residual Connections?

Standard residual connections are one-way: information flows forward only. Evoformer's bidirectional feedback allows:

Deep layers to correct shallow layer errors (Level 1)
Later tokens to inform earlier tokens (Level 2)
Predictions to refine their own inputs (Levels 3-4)
Router and experts to co-evolve (Level 5)

Why Complete Gradient Flow (v1.9.0)?

In prior versions, in-place buffer updates and detached tensors in Evoformer and DualMemory inadvertently blocked gradient flow, meaning some parameters never received training signal. v1.9.0 ensures:

Every parameter has a differentiable path from loss to itself
Feedback loop parameters (revision gates, co-evolve state) are trained effectively
Memory system parameters (state_proj, output_proj) receive gradient signal
ThinkingToggle can be trained end-to-end via sigmoid soft-blending

17. Comparison with Other Architectures

Feature	GPT-4	DeepSeek-V3	Llama 4	Gemma 3	Losion v1.9.0
Architecture	Dense Transformer	MoE + MLA	Dense + iRoPE	Dense + Local/Global	Tri-Jalur Router
Attention	Full	MLA (8× savings)	Interleaved RoPE	Local/Global	MLA + KDA + Lightning + MoBA + Gated + Child-3W
SSM	None	None	None	None	Mamba-2 + Mamba-3 + RWKV-7 + Routing + Liquid + PoST + Structured Sparse + FG2-GDN + DeltaNet
MoE	None	AuxFreeMoE	None	None	AuxFreeMoE + S'MoRE + Symbolic + ∞-MoE + MoHGE + Expert Choice + Gradient Routed + Matryoshka
Router	N/A	Bias-based	N/A	N/A	BiasRouter + ThinkingToggle + Symbolic-MoE
Feedback	None	None	None	None	Evoformer (5 levels)
Memory	KV cache	MLA KV cache	KV cache	KV cache	MLA + Dual Memory (Working + LTM)
RL Training	RLHF	GRPO	RLHF	RLHF	DAPO + GRPO + RLVR
Gradient Flow	Full	Full	Full	Full	Complete (v1.9.0 fix)
Elastic Inference	No	No	No	No	Matryoshka
Speculative Decoding	Yes	Yes	No	No	L-MTP + Mirror + Anchored Diffusion
Expert Prefetching	No	Yes	No	No	Yes (Speculating Experts)

Losion v1.9.0 — Complete Gradient Flow & Vectorized Attention

FilesExpand file tree

ARCHITECTURE.md

Latest commit

History