Elastic Selective Spectral Hybrids for Train-Once, Export-Many Budgeted Inference

Abstract:Deploying a language model under different computing and latency budgets calls for compact models with different quality-cost trade-offs. Towards this end, elastic spectral state space models provide ordered and temporally decomposed channels that can be truncated, but they are associated with linear time-invariant filters that cannot selectively preserve relevant past information or forget irrelevant information as the context evolves. To address this, we introduce the Elastic Selective Spectral Hybrid (ESSH), which realizes each Hankel spectral channel as an independent recurrent unit using fitted damped rotation modes. It also features an input-dependent decay and write/read gates that make temporal retention and state update input-dependent while preserving channel-wise truncation and a structured recurrence for efficient execution. ESSH combines these selective spectral mixers with sliding-window attention and jointly trains multiple capacities by reducing spectral-channel count and feed-forward width at different rates through a two-rate capacity map with full-model distillation. The resulting models support chunked parallel training and fused recurrent decoding while avoiding computation for discarded channels. At full capacity, ESSH achieves language-modeling quality comparable to similarly sized independently trained models, while smaller exports exhibit a smooth quality-cost trade-off. We validate the effectiveness of the proposed framework using language understanding, retrieval, cross-domain text, and DNA experiments, by assessing quality retention and the trade-off against independently trained and elastic baselines. At the 1.53B model configuration, fused batch-one decoding takes 1.37 ms per token on a B300, providing a 2.14-2.80x speedup over the tested Mamba-2 and Mamba-3 implementations and 3.03x over Transformer++ at matched parameter counts.

Machine Learning (cs.LG)

arXiv:2609.32486 [cs.LG]

(or arXiv:2609.32486v1 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2609.32486 Focus to learn more arXiv-issued DOI via DataCite (pending registration)

essh_code/README.md

essh_code/bench/aggregate_matched.py

essh_code/bench/bench_decode_1p5b.py

essh_code/bench/bench_decode_370m.py

essh_code/bench/bench_decode_370m_portable.py

essh_code/bench/bench_matched.py

essh_code/bench/bench_prefill_1p5b.py

essh_code/bench/bench_step_profile_1p5b.py

essh_code/bench/bench_train_throughput.sh

essh_code/bench/cost_timing_recipes.sh

essh_code/bench/decode_fused.py

essh_code/bench/matched_suite_b300/check_environment.py

essh_code/bench/matched_suite_b300/collect.py

essh_code/bench/matched_suite_b300/core.py

essh_code/bench/matched_suite_b300/manifest.json

essh_code/bench/matched_suite_b300/measure.py

essh_code/bench/matched_suite_b300/run.py

essh_code/bench/matched_suite_b300/test_contract.py

essh_code/bench/matched_suite_h200/cfg370.json

essh_code/bench/matched_suite_h200/check_environment.py

essh_code/bench/matched_suite_h200/collect.py

essh_code/bench/matched_suite_h200/core.py

essh_code/bench/matched_suite_h200/manifest.json

essh_code/bench/matched_suite_h200/measure.py

essh_code/bench/matched_suite_h200/measure_small.py

essh_code/bench/matched_suite_h200/run.py

essh_code/bench/matched_suite_h200/test_contract.py

essh_code/bench/sm100_sdpa_shim.py

essh_code/bench/train_rate_probe.py

essh_code/data/merge_llama31_halves.py

essh_code/data/prefetch_eval_corpora.py

essh_code/data/prepare_dna.py

essh_code/data/prepare_fineweb_gpt2.py

essh_code/data/prepare_fineweb_llama31.py

essh_code/data/prepare_genomic_benchmarks.py

essh_code/data/prewarm_datasets.sh

essh_code/data/rebuild_arxiv_pile.py

essh_code/dna/dna_finetune_genomic_benchmarks.py

essh_code/dna/dna_val_loss.py

essh_code/dna/run_dna_finetune_grid.sh

essh_code/essh/act_ckpt.py

essh_code/essh/block_v2.py

essh_code/essh/datainfo.py

essh_code/essh/device_profile.py

essh_code/essh/device_profile.sh

essh_code/essh/distill_chunked.py

essh_code/essh/export_tiers.py

essh_code/essh/fit_hankel_modes.py

essh_code/essh/gsbd.py

essh_code/essh/hankel_fits_permode_L2048_K32_r16.json

essh_code/essh/hankel_fits_permode_L2048_K32_r8.json

essh_code/essh/hankel_fits_permode_L2048_K64_r8.json

essh_code/essh/mix_protocol.py

essh_code/essh/ms_ssd/autotune.py

essh_code/essh/ms_ssd/gate_prod_shape.py

essh_code/essh/ms_ssd/ms_ssd_bwd.py

essh_code/essh/ms_ssd/ms_ssd_bwd_band.py

essh_code/essh/ms_ssd/ms_ssd_fwd.py

essh_code/essh/ms_ssd/ms_ssd_fwd_band.py

essh_code/essh/ms_ssd/test_ms_ssd.py

essh_code/essh/ms_ssd/test_ms_ssd_bwd.py

essh_code/essh/selective_spectral_v2.py

essh_code/eval/count_tier_params.py

essh_code/eval/eval_domain_bpb.py

essh_code/eval/eval_downstream_tasks.py

essh_code/eval/eval_needle_recall.py

essh_code/eval/eval_ppl_exports.py

essh_code/eval/eval_ppl_frontier.py

essh_code/eval/eval_recall_tasks.py

essh_code/eval/fit_comparison.py

essh_code/eval/make_needle_candidates.py

essh_code/eval/probe_mode_horizons.py

essh_code/eval/summarize_needle.py

essh_code/eval/tasks/based/based_drop/default.yaml

essh_code/eval/tasks/based/based_drop/task.py

essh_code/eval/tasks/based/based_fda/default.yaml

essh_code/eval/tasks/based/based_fda/task.py

essh_code/eval/tasks/based/based_nq/nq2048.yaml

essh_code/eval/tasks/based/based_nq/task.py

essh_code/eval/tasks/based/based_squadv2/default.yaml

essh_code/eval/tasks/based/based_squadv2/task.py

essh_code/eval/tasks/based/based_swde/default.yaml

essh_code/eval/tasks/based/based_swde/task.py

essh_code/eval/tasks/based/based_triviaqa/default.yaml

essh_code/eval/tasks/based/based_triviaqa/task.py

essh_code/eval/tasks/check_based_tasks.py

essh_code/posthoc/posthoc_elastify_370m.py

essh_code/refine/shared_refinement.py

essh_code/requirements-mamba.txt

essh_code/requirements.txt

essh_code/scripts/ablation_adaptive_full_share_370m_3b7.sh

essh_code/scripts/ablation_alternative_map_370m_3b7.sh

essh_code/scripts/ablation_channels_kbar16_370m_3b7.sh

essh_code/scripts/ablation_channels_kbar64_370m_3b7.sh

essh_code/scripts/ablation_dff3072_370m_3b7.sh

essh_code/scripts/ablation_dff6144_370m_3b7.sh

essh_code/scripts/ablation_ffn_grad_scale_0p25_370m_3b7.sh

essh_code/scripts/ablation_ffn_ordering_1000_370m_3b7.sh

essh_code/scripts/ablation_flatter_channels_370m_3b7.sh

essh_code/scripts/ablation_flatter_ffn_370m_3b7.sh

essh_code/scripts/ablation_flattest_channels_370m_3b7.sh

essh_code/scripts/ablation_full_share_0p5_370m_3b7.sh

essh_code/scripts/ablation_full_share_0p875_370m_3b7.sh

essh_code/scripts/ablation_full_share_ramp_370m_3b7.sh

essh_code/scripts/ablation_gap_proportional_sampling_370m_3b7.sh

essh_code/scripts/ablation_grad_scale_0p5_both_370m_3b7.sh

essh_code/scripts/ablation_init_all_decays_1e-3_370m_3b7.sh

essh_code/scripts/ablation_init_all_decays_1e-4_370m_3b7.sh

essh_code/scripts/ablation_init_non_oscillatory_370m_3b7.sh

essh_code/scripts/ablation_init_random_modes_370m_3b7.sh

essh_code/scripts/ablation_init_reversed_slots_370m_3b7.sh

essh_code/scripts/ablation_init_s4d_lin_370m_3b7.sh

essh_code/scripts/ablation_init_shuffled_slots_370m_3b7.sh

essh_code/scripts/ablation_kl_coefficient_0p25_370m_3b7.sh

essh_code/scripts/ablation_kl_coefficient_1p0_370m_3b7.sh

essh_code/scripts/ablation_lr_1p5e-3_370m_3b7.sh

essh_code/scripts/ablation_lr_2e-3_370m_3b7.sh

essh_code/scripts/ablation_mixer_only_nesting_370m_3b7.sh

essh_code/scripts/ablation_modes_m16_370m_3b7.sh

essh_code/scripts/ablation_neutral_gates_370m_3b7.sh

essh_code/scripts/ablation_no_ffn_ordering_370m_3b7.sh

essh_code/scripts/ablation_no_pretraining_distillation_370m_3b7.sh

essh_code/scripts/ablation_no_selective_clock_370m_3b7.sh

essh_code/scripts/ablation_progressive_capacity_reduction_370m_3b7.sh

essh_code/scripts/ablation_sampling_five_tiers_370m_3b7.sh

essh_code/scripts/ablation_sampling_ten_tiers_370m_3b7.sh

essh_code/scripts/ablation_steeper_channels_370m_3b7.sh

essh_code/scripts/ablation_weighted_label_1p5ce_370m_3b7.sh

essh_code/scripts/posthoc_recover_essh_dense_370m.sh

essh_code/scripts/posthoc_recover_mamba3_swa_370m.sh

essh_code/scripts/train_dna_essh_0p53m_3b.sh

essh_code/scripts/train_essh_1p5b_100b.sh

essh_code/scripts/train_essh_370m_3b7.sh

essh_code/scripts/train_essh_370m_7b4.sh

essh_code/scripts/train_essh_dense_370m_3b7.sh

essh_code/scripts/train_essh_dense_370m_7b4.sh

essh_code/scripts/train_essh_dense_kbar16_370m_3b7.sh

essh_code/scripts/train_essh_dense_kbar64_370m_3b7.sh

essh_code/scripts/train_gdn_swa_370m_7b4.sh

essh_code/scripts/train_mamba3_swa_370m_7b4.sh

essh_code/scripts/train_matformer_370m_7b4.sh

essh_code/scripts/train_matmamba_370m_7b4.sh

essh_code/scripts/train_specialist_t1_370m_7b4.sh

essh_code/scripts/train_specialist_t4_370m_7b4.sh

essh_code/scripts/train_specialist_t7_370m_7b4.sh

essh_code/scripts/train_tfpp_370m_7b4.sh

essh_code/scripts/train_tfpp_t1_370m_7b4.sh

essh_code/scripts/train_tfpp_t4_370m_7b4.sh

essh_code/scripts/train_tfpp_t7_370m_7b4.sh

essh_code/synthetic/mqar_cells.txt

essh_code/synthetic/mqar_seed_jobs.txt

essh_code/synthetic/run_mqar_cells.sh

essh_code/synthetic/run_mqar_seed_jobs.sh

essh_code/synthetic/run_synth_suite.sh

essh_code/synthetic/synth_tasks.py

essh_code/synthetic/train_synth.py

essh_code/tests/launch_contracts.json

essh_code/tests/test_shell_entrypoints.py

essh_code/tests/test_synthetic_launchers.py

essh_code/tests/verify_derivations.py

essh_code/train/match_dff_params.py

essh_code/train/train_lm.py

(158 additional files not shown) You must enabled JavaScript to view entire file list.

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.