PaperlayerView exact revision on arXiv

Source-linked AI summary

An Exploratory Ablation of a Small MLA--SSM Hybrid Language Model

Christos Koutsiaris

arXiv:2609.29618v1cs.CLcs.AIcs.IRcs.LG

TL;DR

Standard Transformers make local language-model deployment difficult because context growth increases cache, prefill, and memory demands. TALH tests a parallel MLA–SSM hybrid and component variants; in this setup, removing SSM hurt perplexity most, while the dense-FFN variant beat the tested MoE with less peak training memory.

  • Problem

    Standard Transformers challenge local deployment because KV-cache growth, quadratic attention prefill, and total model memory strain commodity hardware.

    Source: §I · p. 1

  • Method

    TALH compares five variants of a hybrid with parallel MLA and recurrent branches under the same optimisation-step and nominal-token budget.

    Source: §I · p. 1, B. Main Training on FineWeb · p. 3

  • Results

    Removing SSM produced the weakest perplexity, while removing MLA left quality close to the full hybrid; dense_ffn had lower loss than the tested MoE and used 3.87 GB less peak training memory.

    Source: §VI · p. 5, §VIII · p. 6

  • Takeaways & Limitations

    The observations support implementation-specific hypotheses, not a general decomposition of MLA and SSM roles or a universal MoE scale threshold.

    Source: §VIII · p. 6

    Validation uses documents from the same two-shard source as training, and missing document identifiers mean train–evaluation overlap cannot be ruled out, limiting absolute perplexity values.

    Source: §VII · p. 6

Abstract

from arXiv · show

We report an exploratory, single-seed ablation of TALH (Adaptive Latent Hybrid), a decoder-only language model with parallel Multi-head Latent Attention (MLA) and a custom recurrent state-space (SSM) branch. Five variants, spanning 117--217M estimated active parameters per token, are trained from scratch on a FineWeb sample for the same number of optimisation steps and tokens. In this specific setup, removing the SSM branch gives the largest degradation in validation perplexity (MLA-only PPL 315), whereas removing MLA has a much smaller effect (SSM-only PPL 239). A dense-FFN hybrid obtains PPL 231, compared with 240 for the tested top-2 ternary-MoE hybrid, while using 3.87 GB less peak training memory. We also preserve a preliminary Apple M3 timing observation: among the five unoptimised implementations, MLA-only has the flattest measured time-to-first-token curve from 512 to 2,048 prompt tokens, although the dense Transformer is much faster in absolute terms. Because the runs are single-seed, parameter counts are unmatched, the evaluation stream may overlap the training source, and raw repeated timing records are unavailable, these results support implementation-specific hypotheses rather than general conclusions about MLA, SSMs, or mixture-of-experts models.

I. INTRODUCTION

The paper explores whether combining MLA, a recurrent SSM branch, and optional sparse expert routing can address language-model quality and local-inference constraints. It tests component contributions and latency under a constrained, single-seed training budget.

I. INTRODUCTION

  • Transformers pose local-deployment challenges because KV-cache growth and quadratic attention prefill burden long contexts, while model memory can exceed commodity hardware.Prior approaches include cache-sharing attention, recurrent SSMs, hybrids, and MLA's low-rank KV compression.Source: §I · p. 1
  • TALH combines MLA with a parallel recurrent branch and optional sparse expert routing to study how these components behave under a constrained budget.The study asks which of MLA or SSM is the primary quality driver, whether MLA cache compression improves Apple Silicon latency scaling, whether ternary MoE helps over dense FFN, and how the hybrid compares with a smaller Transformer.Source: §I · p. 1
  • The five-way FineWeb ablation matches optimisation steps, nominal tokens, data stream, and curriculum, but not FLOPs, parameters, or wall-clock time.The authors describe the comparison as descriptive rather than an estimate of an architecture-only causal effect.Source: §I · p. 1

II. RELATED WORK

Prior work explores attention, SSM, and hybrid designs, while TALH contributes a small-scale, from-scratch ablation of hybrid components. Its results are framed as implementation-specific rather than a general scale threshold.

II. RELATED WORK

  • Earlier work reports linear-memory SSMs and hybrid SSM-attention models with perplexity or inference advantages, while later hybrids combine these components in varied layer or head designs.The cited examples include S4/Mamba, H3, Jamba, Samba, Hymba, and Zebra-Llama.Source: §II · pp. 1–2
  • Hymba and Zebra-Llama are close design precedents, but TALH trains from scratch, directly ablates MoE against dense FFN, and adds preliminary Apple Silicon timing.The paper presents its contribution as a transparent ablation and negative result at smaller scale.Source: §II · p. 2
  • MoE gains at large scales do not establish that this smaller ternary configuration should help; outcomes depend on routing, capacity, precision, and comparison budget.The paper notes compact MoE gains around 200M active parameters under different designs and controls.Source: §II · p. 2

III. ARCHITECTURE

  • TALH is described as having three architectural components.The supplied architecture-section passage does not specify those components further.Source: §III · p. 2

A. Multi-head Latent Attention (MLA)

TALH pairs latent-cache attention with a recurrent branch and compares a ternary expert feed-forward block against a dense alternative. The supplied passages describe their mechanisms and intended scope.

A. Multi-head Latent Attention (MLA)

  • MLA stores a low-dimensional latent instead of full key and value vectors, yielding a theoretical 4× per-head compression ratio in this setup and an 8× total-cache reduction in preliminary 15M experiments.The latent dimension is 200 versus d_model 800, and inference projects the latent back into key and value space.Source: A. Multi-head Latent Attention (MLA) · p. 2

B. State Space Model (SSM) Branch

  • The custom selective recurrence, unlike Mamba or Mamba-2, has fixed-size state and linear sequence-length arithmetic; in TALH it runs parallel to MLA and a learned projection with sigmoid gate combines their outputs.The implementation is intended for architectural exploration rather than as a production SSM kernel.Source: B. State Space Model (SSM) Branch · p. 2

C. Ternary Mixture-of-Experts (MoE)

  • The full variant uses top-2 routing among eight ternary-weight experts, while dense_ffn uses a standard BF16 SwiGLU layer and has the lowest loss in this ablation.The ternary weights are quantised from FP32 master weights during the forward pass; the masters are not packed into a 1.58-bit representation.Source: C. Ternary Mixture-of-Experts (MoE) · p. 2

D. Ablation Variants and Parameter Counts

The study estimates active parameters as a parameter-access convention, not a FLOP estimate, and flags a shared FFN-width choice as a possible confound. A smaller pilot reports cache, memory, and perplexity observations.

D. Ablation Variants and Parameter Counts

  • Active-parameter estimates count weights touched per token, not FLOPs, and do not account for the different execution costs of shared weights and the recurrent branch.For MoE variants, the estimate includes non-expert weights and two of eight experts per layer.Source: D. Ablation Variants and Parameter… · p. 2
  • All variants use d_ff = 1,600, a 2× hidden-width ratio rather than the often-used 4×; widening the baseline could raise it toward 163M parameters and reduce the quality gap.The same width need not affect different architectures equally.Source: D. Ablation Variants and Parameter… · p. 2

A. Preliminary Experiments at 15M Parameters

B. Main Training on FineWeb

The study trains five variants on FineWeb with matched steps, nominal tokens, source stream, and curriculum, but not matched parameter counts, operations, or wall-clock time. Evaluation uses a stream that may overlap the training source, and only two hybrid variants continue into Phase 2.

B. Main Training on FineWeb

  • The evaluation stream comes from the same FineWeb shards as training, and the released logs do not establish that it is disjoint from training documents.The implementation labels this stream “validation,” but the paper does not claim a clean held-out estimate.Source: B. Main Training on FineWeb · p. 3
  • Training uses one A100-SXM4-40GB GPU and a context curriculum of 256 tokens for steps 0–2,999, 512 for steps 3,000–4,999, and 1,024 for steps 5,000–7,000.Phase 1 covers the 256- and 512-token stages.Source: B. Main Training on FineWeb · p. 3
  • At step 4,000, all five variants share the same optimisation steps and nominal tokens, but parameter counts, operations, and wall-clock time are unmatched.This is the primary comparison point; Fig. 3 reports step- and token-matched quality at that step.Source: B. Main Training on FineWeb · p. 3, Figure 3 · p. 3
  • Only full and dense_ffn continue to 7,000 steps and the 1,024-token stage; their Phase 2 results are not directly comparable with Phase 1-only variants.The other three variants stop after 4,000 steps because of GPU budget constraints.Source: B. Main Training on FineWeb · p. 3
  • Validation perplexity is evaluated at sequence lengths 512 and 1,024, with one training run per variant and no multi-seed replication.Training uses AdamW with cosine decay, an initial learning rate of 10^-4, 400 warm-up steps, gradient clipping at 0.5, and BF16 mixed precision.Source: B. Main Training on FineWeb · p. 3

C. Local Inference Benchmark

At the step- and token-matched Phase 1 comparison, dense_ffn has lower measured perplexity than the dense Transformer baseline and the tested MoE hybrid. The supplied inference passage describes preliminary TTFT measurement conditions and reproducibility limits, but does not provide timing outcomes.

C. Local Inference Benchmark

  • TTFT was measured at 512, 1,024, and 2,048 prompt tokens on an unoptimised MacBook M3 prototype, but raw per-trial timing records are unavailable.The passage cautions that absolute latency should not be compared directly with production-optimised inference frameworks.Source: C. Local Inference Benchmark · p. 3

A. Step- and Token-Matched Quality: Phase 1

  • The dense_ffn and full hybrids lead the Phase 1 validation-loss curves throughout, while mla_only consistently trails.Fig. 4 plots validation loss for all five variants.Source: Figure 4 · p. 4
  • PPL 231 for dense_ffn versus 263 for transformer_dense, a 12% reduction under the same step and nominal token budget.Table II reports Phase 1 results at step 4,000, with PPL measured at 512-token context.Source: A. Step- and Token-Matched Quality:… · p. 3, Table II · p. 4
  • The dense_ffn comparison is not replicated across seeds and is not an architecture-only or statistically tested effect.The supplied passage notes that the hybrid has more active parameters and used substantially more wall-clock time.Source: A. Step- and Token-Matched Quality:… · p. 4
  • PPL 315 for mla_only, compared with PPL 239 for ssm_only, PPL 240 for full, and PPL 263 for transformer_dense.Within this run and implementation, removing the recurrent branch corresponds to the largest quality degradation; the result does not establish the same role decomposition for other SSM or MLA designs.Source: A. Step- and Token-Matched Quality:… · p. 4
  • PPL 231 for dense_ffn versus PPL 240 for full with MoE, with dense_ffn using 3.87 GB less peak training memory.The authors did not measure routing specialisation, so this experiment does not identify why the tested MoE configuration underperforms.Source: A. Step- and Token-Matched Quality:… · p. 4

B. Extended Training: full and dense_ffn

In extended training, dense_ffn retains its validation-loss advantage over full: it matches full’s loss about 1,000 steps earlier and has lower loss at step 7,000. The 1,024-token stage required smaller batches, and full used more peak memory.

B. Extended Training: full and dense_ffn

C. Quality vs. Memory Trade-off

At step 4,000, dense_ffn has the lowest measured perplexity and peak memory among the hybrid variants, while transformer_dense uses less memory but has higher measured perplexity. In preliminary M3 measurements, mla_only has the flattest observed TTFT curve, although transformer_dense is much faster in absolute terms.

C. Quality vs. Memory Trade-off

  • Transformer_dense uses 5.36 GB but has 14% higher measured perplexity; the roughly 19 GB gap cannot be assigned to useful architectural capacity alone.The gap also reflects different active parameter counts, optimiser state, the recurrent branch, MLA projections, and gating.Source: C. Quality vs. Memory Trade-off · p. 4
  • Among the three hybrid variants at step 4,000, dense_ffn has the lowest measured perplexity and peak memory.Fig. 6 plots perplexity against peak GPU training memory for all five variants.Source: C. Quality vs. Memory Trade-off · p. 4, Figure 6 · p. 5

D. Local Inference Latency on Apple Silicon

  • From 512 to 2,048 tokens, mla_only TTFT grows 1.32×, versus 5.5× for transformer_dense, 3.3× for dense_ffn, and 2.7× for full.These are preliminary observations from a single run, with no repetitions or uncertainty estimates.Source: D. Local Inference Latency on Apple… · p. 4, Figure 7 · p. 5, Table IV · p. 5
  • The TTFT pattern establishes neither asymptotic complexity nor KV-cache compression as its cause, and is only a hypothesis for a future controlled benchmark.Repetitions, variance, and inference-memory measurements are unavailable.Source: D. Local Inference Latency on Apple… · p. 5
  • At 512 tokens, transformer_dense takes 114 ms versus 2,560 ms for dense_ffn; the comparison reflects implementation differences, including fused attention versus separate unfused operations.The measurements use an unoptimised prototype stack.Source: D. Local Inference Latency on Apple… · p. 5

VI. DISCUSSION

The discussion treats the component ablation and latency findings as hypotheses about this implementation, not general architectural conclusions. It also notes that the tested MoE result has several unresolved explanations and that comparisons with larger prior systems differ in training setup.

VI. DISCUSSION

  • Removing SSM produces the worst perplexity, while removing MLA leaves quality close to the full hybrid; single-seed runs and unmatched parameter counts prevent a stronger causal claim.The study also lacks a dense-FFN component ablation.Source: §VI · p. 5
  • The experiment cannot distinguish whether the tested MoE underperforms because of ternary forward weights, router behaviour, expert capacity, or parameter allocation.Routing statistics and a dense-expert MoE control are needed before attributing the result to specialisation; FP32 master weights and optimiser states for all experts remain resident during training.Source: §VI · p. 5
  • Hymba [6] and Zebra-Llama [7] report stronger quantitative results, but use substantially more tokens and production infrastructure or fine-tune an already-trained base model.TALH trains from scratch with a constrained budget and is presented as an ablation study, not a state-of-the-art system.Source: §VI · p. 5
  • The 114 ms versus 2,560 ms TTFT gap at 512 tokens is evidence about the current prototype, not that hybrid architecture is unsuitable for edge deployment.The passage attributes the gap to implementation optimization differences, including fused, hardware-aware attention kernels for standard MHA.Source: §VI · p. 5

VII. LIMITATIONS

The study’s interpretation is limited by single runs, parameter and training-cost mismatches, an evaluation stream that may overlap training, and incomplete inference benchmarking. Its FineWeb perplexity results also do not directly assess downstream reasoning or instruction following.

VII. LIMITATIONS

  • All variants were trained once, so perplexity differences—especially the Phase 1 values 231, 239, and 240—may not be statistically robust without multi-seed replication.The values correspond to dense_ffn, ssm_only, and full, respectively.Source: §VII · p. 5
  • Active parameter counts are imbalanced: dense_ffn has 171M versus 117M for transformer_dense, while full, ssm_only, and mla_only have about 217M, 203M, and 177M.The 12% quality advantage over the Transformer baseline is therefore not purely architectural; all variants also use dff = 1,600, and the passage says the gap would likely be smaller with a standard 4× FFN ratio.Source: §VII · pp. 5–6
  • Phase 1 matches steps and nominal tokens, not FLOPs, parameters, or wall-clock time; Phase 2 extends only dense_ffn and full, so those results cannot be compared with the other variants.The reported 4,000-step times range from approximately 0.58 h for transformer_dense to 19.3–20.5 h for recurrent hybrids.Source: §VII · p. 6
  • Training–evaluation overlap cannot be ruled out because the FineWeb evaluation stream skips 10,000 documents from the same two-shard source and consumed document identifiers are not archived.This limits absolute perplexity values, although all variants use the same procedure.Source: §VII · p. 6
  • No downstream task evaluation is reported; FineWeb perplexity does not directly measure reasoning or instruction-following ability..Source: §VII · p. 6
  • Inference benchmarking omits decode throughput, peak inference memory, and power, while unavailable raw repetitions and software versions prevent assessing variance and exact reproduction.The prototype uses no quantisation, kernel fusion, or prefix caching, and production-framework comparisons would be inappropriate.Source: §VII · p. 6

VIII. CONCLUSION

In this exploratory ablation, removing the recurrent branch produces the weakest perplexity, while MLA-only has the flattest preliminary TTFT growth. The dense-FFN hybrid also outperforms the tested ternary-MoE variant on loss while using less peak training memory.

VIII. CONCLUSION

  • Removing the recurrent branch produces the weakest perplexity in the study, while removing MLA leaves quality close to the full ternary-MoE hybrid.These are hypotheses about the tested implementation, not a general role decomposition across MLA–SSM architectures.Source: §VIII · p. 6
  • 1.32× TTFT growth for a fourfold context increase (512 to 2,048 tokens) was recorded for MLA-only, versus 2.7–5.5× for the other variants.Without repeated trials, memory measurements, or a prefill/decode decomposition, the observation is preliminary and does not confirm complexity or causality.Source: §VIII · p. 6
  • 3.87 GB less peak training memory accompanies lower loss for the dense feed-forward hybrid than for the tested ternary-MoE variant.The authors frame this as a negative result for that configuration, not evidence of a universal MoE scale threshold.Source: §VIII · p. 6
  • Disjoint-data and multi-seed replication, parameter/FLOP-matched controls, routing diagnostics, and repeated prefill, decode, and memory benchmarks are identified as next steps.The proposed benchmarks include architecture-aware checkpoint conversion.Source: §VIII · p. 6
Loading 2609.29618v1…