ChimeraPKM-300M

Comparative Benchmark Report: ChimeraPKMv2 vs. Industrial Baselines

We present ChimeraPKMv2, a fully attention-free hybrid language model combining a gated linear State Space backbone with differentiable Product Key Memory. Our architecture demonstrates constant recurrent-state memory usage during inference and training while preserving competitive factual retrieval through explicit associative memory.

This benchmark report presents a rigorous evaluation of the ChimeraPKMv2 architecture (300M parameters) against two primary baselines: a Decoder-only Transformer (with causal self-attention and KV Cache) and a Pure SSM (Gated Linear Scan without external memory) of equivalent size.

All tests were conducted natively on Apple Silicon GPU (Metal) using the Apple MLX framework.

---

1. Complete Architectural Specifications

ChimeraPKMv2 is a hybrid, non-attention-based language model. It replaces Multi-Head Self-Attention with a gated linear recurrent operator (ChimeraMixer) and integrates a differentiable 2D key-value retrieval layer (PKMRetrievalV2).

The model is structured in a symmetric sandwich layout: 4 recurrent blocks -> 1 PKM Memory layer -> 4 recurrent blocks.

Input Tokens ---> [Embedding] 
                    |
             [Blocks 1 - 4]  (ChimeraMixer recurrent scan + MLP)
                    |
            [PKMRetrievalV2] (2D Hard routing + Differentiable STE Lookup)
                    |
             [Blocks 5 - 8]  (ChimeraMixer recurrent scan + MLP)
                    |
               [RMSNorm] ---> [LM Head] ---> Logits

---

1.1 The Recurrent Core: ChimeraMixer

The ChimeraMixer is a gated linear State Space Model (SSM) variant optimized for parallel prefix scans and high-throughput recurrent generation. Unlike attention which compares all token pairs O(T^2), ChimeraMixer propagates sequence history via a first-order recurrence relation.

Mathematical Formulation

Given an input sequence X in R^{B x T x D} where B is the batch size, T is the sequence length, and D is the embedding dimension:

1. Short Convolution: X is processed by a group-convolution with kernel size K=4 to enforce local spatial context and prevent token order vanishing: $$X_{conv} = \text{Conv1d}(X, \text{kernel}=4, \text{groups}=D)$$

2. Projection Layers: We compute Query (Q), Key (K), Value (V), Decay (a), and Input Gate (b) projections: $$Q = W_q X_{conv}, \quad K = \text{Norm}(W_k X_{conv}), \quad V = W_v X_{conv}$$ $$a = 0.9 + 0.1 \cdot \sigma(W_{gate} X_{conv})$$ $$b = \sigma(W_{beta} X_{conv}) \cdot \sigma(W_{salience} X_{conv})$$ Where \sigma is the sigmoid function. The biases for the gate and salience projections are initialized to 4.0 to ensure the recurrent channel gates are open at initialization.

3. Gated Delta Scan (Recurrence): For each head h, the state S_t updates according to the gated recurrence: $$S_t = a_t S_{t-1} + b_t (K_t \otimes V_t)$$ $$Y_t = Q_t S_t$$ To train efficiently on GPU, this recurrence is computed in parallel chunks of size C=16 using the associative prefix scan. The matrix inversion for parallel chunk solving is stabilized using a triangular Neumann series approximation (_unit_lower_inv): $$I - A^{-1} \approx \sum_{k=0}^{\lceil \log_2 C \rceil} (-A)^{2^k}$$

4. Output Projection: The output of the heads is concatenated and projected back to the hidden dimension: $$\text{ChimeraMixer}(X) = W_o Y$$

---

1.2 The Associative Memory Core: PKMRetrievalV2

The PKMRetrievalV2 layer acts as a differential lookup table containing 16,384 memory buckets (size 128 x 128). It allows the model to retrieve static factual embeddings in O(1) time without keeping facts inside the recurrent states of the Mixer.

1. Differentiable 2D Keys Routing

Instead of matching a query against 16k keys (which is computationally expensive), we project the hidden query x into two sub-keys of size D_key = 256: $$q = W_{query} x, \quad q_x = \text{LayerNorm}(q_{1:D_{key}}), \quad q_y = \text{LayerNorm}(q_{D_{key}:2D_{key}})$$ These sub-keys are matched independently against two small key-matrices C_x in R^{128 x D_key} and C_y in R^{128 x D_key}: $$s_x = q_x C_x^T, \quad s_y = q_y C_y^T$$ $$p_x = \text{Softmax}(s_x), \quad p_y = \text{Softmax}(s_y)$$

2. Straight-Through Estimator (STE)

During training, we want the routing to be sparse (hard selection of a single bucket to save computation), but sparse argmax is non-differentiable. We use the Straight-Through Estimator to route gradients: $$i_x = \text{Argmax}(s_x), \quad i_y = \text{Argmax}(s_y)$$ $$M_{hard} = \text{Buckets}[i_x, i_y]$$ $$M_{soft} = \sum_{j} \sum_{k} p_x[j] \cdot p_y[k] \cdot \text{Buckets}[j, k]$$ $$ ext{Memory} = M_{soft} + \text{stop_gradient}(M_{hard} - M_{soft})$$ This allows the forward pass to use the exact hard-routed embedding, while the backward pass flows gradients smoothly through the soft probabilities p_x and p_y to update the key projections.

3. Confidence-Aware Gating

To prevent the model from retrieving memories when it is not confident about the lookup, we scale the memory injection using a confidence gate. The confidence is the joint probability of the top-1 keys: $$ ext{Confidence} = \max(p_x) \cdot \max(p_y)$$ $$ ext{Gate} = \sigma(W_{inject} x + \lambda_{pkm} \cdot \text{Confidence})$$ $$ ext{Output} = x + \text{Gate} \cdot W_{out}( ext{Memory})$$ Where \lambda_{pkm} = 2.0.

4. Orthogonality Loss

To ensure that the x and y key projections learn independent dimensions (preventing the 2D grid from collapsing into a 1D diagonal line), we penalize the cosine similarity between q_x and q_y: $$\mathcal{L}{ortho} = \frac{1}{B \cdot T} \sum{b, t} \left| \frac{q_x \cdot q_y}{|q_x| |q_y|} \right|$$ This loss is scaled by a = 0.001 and added to the training objective.

---

2. Comparative Benchmarks

2.1 Incremental VRAM Allocation (Post Warm-up)

Measures the newly allocated VRAM required as context length scales from 256 to 4096 tokens during Inference (Forward Pass) after initial compilation.

[!NOTE] Values represent incremental recurrent-state allocations post-compilation, not total memory usage. Embedding matrices, MLP activations, and LayerNorm activations are always present and scale linearly with sequence length. The key distinction is that the Transformer KV Cache grows linearly $O(T)$ with context length (adding hundreds of MB at 4K tokens), whereas SSM and ChimeraPKMv2 exhibit no sequence-dependent recurrent allocation growth after initial graph compilation. Time complexity remains $O(T)$ for all models.

---

2.2 Generation Throughput (Tokens / Second)

Measures parallel forward pass throughput across different sequence lengths.

---

2.3 Incremental Training Memory Footprint (Backprop VRAM Delta in MB)

Measures the VRAM allocated for saving activation graphs during Training (Forward + Backward Pass) after initial compilation.

[!IMPORTANT] Because backpropagation in Self-Attention requires caching $Q \times K^T$ activation maps that grow quadratically, fine-tuning standard Transformers on long contexts is extremely expensive. ChimeraPKMv2 exhibits no sequence-dependent recurrent-state activation growth, significantly reducing peak VRAM overhead during backpropagation. Note: total memory is not zero — embedding tables, MLP, and LayerNorm activations are always present. The saving is specifically in the recurrent-state component.

---

2.4 Quantization Sensitivity & Validation

Evaluates weight matrix degradation when compressing parameters to lower bit depths (8-bit and 4-bit integers).

[!TIP] Under fake 4-bit quantization of linear weight matrices, the domain-specific validation perplexity shows negligible degradation ($5.285 \rightarrow 5.285$), indicating that PKM routing coordinates appear numerically robust at this precision level. This result should be interpreted with caution: perplexity on a small domain corpus may not capture degradation on general downstream tasks (MMLU, GSM8K, HumanEval). Full INT4 validation requires evaluation on standardized public benchmarks.

---

3. Empirical Quality Validation & Ablation Study

To prove the specific contribution of the Product Key Memory (PKM) to the model's factual capacity, we conduct an ablation study comparing the full model against its components on a specialized Q&A corpus.

3.1 Corpus & Dataset Characteristics

The evaluation dataset is a domain-specific Q&A dataset containing:

• Total Pairs: 2,902 Q&A pairs (90% training / 10% validation).
• Token Count: 348,040 tokens.
• Domain Breakdown: Python programming syntax (10%), Macroeconomics principles (45%), and Discrete Math/Propositional Logic (45%).
• Format: Structured prompt wrappers (Domanda: [Text] Risposta: [Text] <|endoftext|>).

[!IMPORTANT] This dataset is highly domain-specific and serves as a Proof-of-Concept (PoC) demonstration rather than a general-purpose pre-training validation. It is designed to verify factual routing under strict parameter alignment on a local budget.

3.2 Experimental Controls

To guarantee strict comparability, all ablated models were initialized with the same random seed and trained under identical hyperparameters:

• Learning Rate (LR): $5 \times 10^{-5}$
• Sequence Length: 256
• Optimization Steps: 3,000 steps
• Effective Batch Size: 8 (gradient accumulation factor of 4 with physical batch size of 2)

3.3 Loss & Perplexity (PPL) Ablation Table

Validation metrics were evaluated across 3 independent seeds to establish standard deviation bounds ($\pm$).

[!IMPORTANT] The validation split is strictly decoupled to prevent data leakage. However, given the specialized nature and small scale of the corpus (2,902 Q&A pairs), a degree of domain overfitting is expected. Scaling pre-training to diverse web-scale data is necessary to evaluate the model's out-of-distribution (OOD) generalization capabilities.

Integrating the PKM layer improves validation perplexity from 24.58 to 14.21 (a 42.1% reduction). Unfreezing the pre-PKM encoder layers is necessary to prevent feature representation collapse, allowing queries to route to distinct memory coordinates, raising Exact Match (EM) accuracy to 96.4%.

---

4. Product Key Memory Routing Dynamics

To verify the routing behavior of the $128 \times 128$ Product Key grid (16,384 total buckets) under active backpropagation:

• Active Unique Bucket Rate: During Q&A training, the unique active bucket rate reaches 20% - 25% (approx. 3,200 - 4,000 unique coordinates), demonstrating distributed routing.
• Routing Entropy & Gini Coefficient: Without the Orthogonality Loss penalty, the sub-keys $q_x$ and $q_y$ degenerate to the diagonal ($x = y$). The Gini coefficient under collapsed diagonal routing is $\approx 0.95$ (highly unequal, dead buckets). Applying the Orthogonality Loss reduces the Gini coefficient to $0.38$, showing uniform routing coverage.

---

5. Advanced Offline Diagnostics & Interpretability

To analyze model behavior and identify current limitations, we conducted four offline scientific tests on Apple Silicon without external compute. All tests use the production weights (chimera_pkm_surgical_best.safetensors, 191 tensors loaded).

---

5.1 PKM Compositionality Test

Goal: Verify whether the PKM has learned semantic representations or is functioning as a rigid textual dictionary. The model is queried with paraphrased versions of corpus questions that were never seen during training.

Result: 0/4 (0%) compositionality. The PKM is functioning as a textual dictionary rather than a semantic memory. Paraphrased queries activate incorrect routing coordinates, retrieving unrelated domain content. This is consistent with the small corpus size (348k tokens) and is a primary motivation for web-scale pre-training.

---

5.2 Variable-Depth Needle in a Haystack

Goal: Map the SSM's recurrent memory retention curve by inserting a target fact at different positions within a 1,231-token context.

Result: 0/5 at all depths. The model produces the same response regardless of needle position, indicating it ignores the injected context entirely and retrieves a learned pattern from training. This is consistent with known SSM behavior at the 300M parameter scale with limited pre-training. Scaling model size and training budget is expected to improve context retention.

---

5.3 Bucket Activation Map (128×128 Grid)

Goal: Measure the actual routing distribution of the PKM during inference on the validation split. Raw heatmap saved in results/bucket_heatmap.csv and results/bucket_heatmap.png.

Top-10 most activated buckets:

[!IMPORTANT] A Gini of 0.998 with only 0.5% active buckets confirms that the PKM is severely under-utilized on this corpus. The Orthogonality Loss helps prevent diagonal collapse but is insufficient at this data scale to force distributed routing across the 16,384 bucket grid. This establishes a critical baseline: the target metric for web-scale pre-training is a Gini below 0.5 with at least 20–25% active buckets (as observed during training, 3,200–4,000 unique coordinates).

---

5.4 Out-of-Distribution (OOD) Perplexity

Goal: Quantify domain overfitting by comparing perplexity on the training corpus versus unseen Italian text from different domains.

[!IMPORTANT] The OOD perplexity ratios (×5.7 to ×29.2) confirm significant domain specialization. Notably, the lowest OOD ratios are on factual/news-style Italian text (calcio, Roma), while the highest are on instructional/procedural text (recipe). This suggests the model has learned the Q&A question-answer structure of the corpus more than the underlying domain knowledge. Web-scale pre-training on diverse Italian and multilingual text is necessary before general-purpose inference.

---

5.5 Summary: Diagnostic Conclusions

[!NOTE] These results are entirely consistent with the PoC scope of this work (348k training tokens, single specialized domain). They do not invalidate the architectural hypothesis; rather, they establish precise quantitative targets for the next research phase: web-scale pre-training, larger model variants, and multi-domain corpora. The Gini coefficient (0.998 → target < 0.5) and OOD ratio (×29 → target < 3) serve as measurable milestones for future experiments.

---

6. Task Suitability Guidelines

---

7. Future Academic Roadmap & Known Limitations

To prepare this architecture for submission to conferences (such as EMNLP, ICLR, or NeurIPS), the following research steps are being prioritized:

1. Modern Baselines Comparison: We plan to benchmark ChimeraPKMv2 directly against modern linear-time models such as Mamba, RWKV-6, Griffin, and DeltaNet of equivalent parameters on the same corpus and compute budget.
2. Downstream Task Evaluations: Quantifying generation quality under 4-bit and 8-bit quantization using standard evaluation suites (MMLU, GSM8K, HumanEval, ARC) instead of relying solely on domain perplexity or MSE.
3. Needle in a Haystack (Context Retention): Evaluating context retrieval accuracy across context lengths up to 128k tokens to precisely map the SSM's information retention decay threshold. The current 300M PoC model fails at 2K tokens for casual out-of-distribution facts, consistent with typical SSM behavior at this scale.
4. Scaling Laws of Associative Memory: Investigating how perplexity scales when expanding the grid size from 16,384 ($128 \times 128$) to 65,536 ($256 \times 256$) buckets, and whether the Gini routing uniformity is preserved at scale.
5. Web-Scale Pre-Training: Validating that the PKM contribution (PPL: 24.58 → 14.21, EM: 72% → 96%) generalizes beyond the current 348k-token PoC corpus to diverse web-scale data (e.g. RedPajama, Pile).

[!WARNING] The authors explicitly note that the SSM+PKM hybrid concept has precedents in the research literature (e.g. memory-augmented SSMs, external memory networks). The contribution of this work lies in the specific integration: a fully attention-free sandwich architecture combining ChimeraMixer (Neumann-stabilized parallel scan), 2D STE routing, confidence-aware gating, and orthogonality regularization — validated under strict experimental controls on Apple Silicon.