{"id":"ff8ba055-632c-463e-bf80-de67f80b6eaa","arxiv_id":"2504.20922","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"DYNAMAX adds early-exit classifiers to Mamba and Transformer LLMs and uses Mamba blocks as classifiers, showing computational savings with limited performance loss.","lead":"This paper introduces early exit mechanisms for Mamba-based language models, letting inference stop early when the model is confident, and also uses Mamba blocks as the exit classifiers themselves. It reports computational savings with modest accuracy loss on several question-answering benchmarks.","discovery_kind":"new_application","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Unquantified post-hoc exclusion of degenerate low-threshold configurations inflates the reported savings-accuracy trade-off.","rationale":"I focused on the explicit data-exclusion rule in Section V because it directly controls the quantitative evidence for the central claim of savings with limited degradation. The reader's weakest_assumption (Mistral vs Codestral confound) also weakens generalization, but the exclusion rule is more load-bearing: it can change the plotted Pareto curves themselves, independent of model choice. The reader listed this issue in the rationale, so agreement is partial. A re-analysis that reports all configurations would settle it. I do not call for rejection; the framework and code release are real contributions, and the issue is addressable. Hence the CONDITIONAL verdict stands (UNCHANGED).","tokens_in":11182,"tokens_out":5631,"duration_ms":61075,"concrete_test":"Re-run all threshold configurations from Section V without excluding repeated-token degeneracies; report exact match, F1, BLEU, and Rouge for every threshold, and add a column with the number of degenerate generations per configuration. Then recompute the Pareto curves. If the omitted low-threshold points fall far below the shown curves, the claimed savings-accuracy trade-off does not hold. Also report the exclusion criterion and the count of excluded points for each model and classifier.","verdict_should_be":"UNCHANGED","load_bearing_attack":"Section V states: 'Sometimes, when the threshold is set too low, the model proceeds to output inconsistent responses that are composed of the same repeated token. If this event occurs a certain number of times, that configuration for the evaluation is not considered valid and then not shown in the graphs.' The exclusion rule is unspecified ('a certain number of times'), no count of excluded configurations is reported, and no sensitivity analysis is given. Because the x-axis of every EE plot is the threshold sweep, removing failed low-threshold points selects the Pareto frontier and directly supports the headline claim of 'computational savings with limited performance degradation.' The rule applies only to EE configurations, not to layer-pruning baselines, so the EE-versus-pruning comparison is also affected. If the excluded degenerate outputs were counted as failures, the accuracy at high savings could drop substantially. This is the load-bearing weak point: the quantitative evidence for the central claim relies on curves that may be cherry-picked by an unstated rule.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper introduces DYNAMAX, a framework that adds early exit (EE) mechanisms to decoder-only Transformers and Mamba-based LLMs, and also repurposes Mamba blocks as lightweight EE classifiers for both Mamba and Transformer backbones. The classifiers are trained by knowledge distillation from the full model, with a top-k relaxation of the oracle target to improve training stability. During inference, tokens may exit early when a classifier's confidence exceeds a threshold; for Transformers, the paper proposes copying KV caches from the exit layer to subsequent blocks, while for Mamba it considers either partial recomputation or skipping state updates. Experiments compare Mistral 7B (Transformer) with Codestral Mamba 7B on TriviaQA, CoQA, and TruthfulQA, measuring accuracy versus computational savings, and compare EE against layer pruning. The authors claim this is the first framework to integrate EEs into Mamba architectures and to use Mamba as an EE classifier, and report that Mamba-based classifiers provide a favorable savings-accuracy trade-off.","tokens_in":11397,"tokens_out":2712,"duration_ms":30205,"significance":"If the central claims hold, the paper makes a useful contribution by extending early-exit techniques to state-space models, which are increasingly important for efficient long-context inference, and by demonstrating that Mamba blocks can serve as effective confidence predictors. The release of source code and the use of public models and datasets support reproducibility and facilitate follow-up work. However, the empirical support for the headline claim is currently weakened by an unquantified post-hoc exclusion rule for degenerate low-threshold configurations, a confounded comparison between two models that differ in pretraining data and specialization, and an incompletely specified KV-cache copying mechanism. These issues are load-bearing because the paper's quantitative conclusions rest on the threshold-sweep curves and on the architectural comparison.","major_comments":[{"comment":"The exclusion rule for degenerate outputs is not specified quantitatively: the paper states 'If this event occurs a certain number of times, that configuration for the evaluation is not considered valid and then not shown in the graphs,' but it does not give the number, the total count of excluded configurations, or the number of degenerate sequences per configuration. Since the x-axis of every EE plot is the confidence-threshold sweep, removing low-threshold degenerate points preferentially trims the high-savings end of the curve, which directly inflates the reported savings-accuracy trade-off. Moreover, the rule is applied only to EE configurations and not to layer-pruning baselines, biasing the comparison. The authors should report the exclusion criterion, the number of excluded configurations per task and model, and a sensitivity analysis in which degenerate outputs are counted as failures; the figures should be re-plotted accordingly.","section":"Section V"},{"comment":"The architectural comparison is confounded by pretraining: Mistral 7B is a general-purpose Transformer, whereas Codestral Mamba 7B is 'specifically optimized for code tasks' (Section IV-A). Any observed difference in early-exit behavior between the two backbones could be due to the different pretraining corpora and domains rather than to the architectural properties under study. The authors should either use Mamba and Transformer models trained on the same data (e.g., the same family with comparable checkpoints), or explicitly argue why the pretraining difference is irrelevant to early-exit behavior, or provide a matched ablation that isolates architecture from pretraining.","section":"Section IV-A and V"},{"comment":"The proposed KV-cache copying trick is incompletely specified and unvalidated: the passage describing the update of subsequent Transformer blocks ends mid-sentence ('will have its KV cache updated with the keys and value from the'), and no experiment isolates the effect of copying versus recomputation on generation quality. The claim that copying cached representations from the exit layer preserves enough information for correct generation is an axiom that needs empirical support. Please complete the algorithm description, state exactly which keys and values are copied, and provide an ablation comparing (a) copying, (b) partial forward recomputation, and (c) no state update, with respect to accuracy and consistency.","section":"Section IV-D"},{"comment":"All experimental results are presented as single deterministic-looking curves without error bars, repeated runs, or statistical significance tests. The training of EE classifiers involves stochastic optimization, and threshold sweeps produce curves whose differences between EE configurations and layer pruning may be within run-to-run variance. Adding multiple seeds (or bootstrap confidence intervals over the evaluation samples) is necessary to support the qualitative claims such as 'Mamba shows as a better model for EE prediction' and 'EE configurations provide better performance.'","section":"Figures 1–5"}],"minor_comments":[{"comment":"In Algorithm 1, the condition 'exit output ≥ θ' is ambiguous because the classifier outputs two values passed through a softmax; clarify that the exit decision uses the probability of the 'exit' class, and define the notation for 'partial forward(z)'.","section":"Algorithm 1"},{"comment":"The sentence describing the KV-cache copying mechanism is syntactically incomplete (it ends with 'the keys and value from the'); please finish the sentence and provide a precise description of what is copied and how the cache tensors are sized.","section":"Section IV-D"},{"comment":"The phrase 'a larger spawn over the x axis' appears to be a typo; 'span' is likely intended. Also, the phrase 'a certain number of times' should be replaced with a concrete numeric criterion.","section":"Section V"},{"comment":"The paper states that a configuration of four early exits is placed in the second half of the model's backbone, but it does not specify the exact layer indices for either Mistral or Codestral Mamba. Providing these indices would improve reproducibility.","section":"Section V-B"},{"comment":"The repetition penalty is mentioned in the Mamba inference paragraph but its value and how it is applied are not described; please give the formula or reference.","section":"Section IV-B"},{"comment":"The related-work discussion would benefit from a direct quantitative comparison with existing EE frameworks for decoder-only LLMs (e.g., CALM and EE-LLM) rather than only citing them; at minimum, the chosen baselines (layer pruning and full model) should be supplemented with a CALM-style baseline on the same architecture.","section":"Related Work"}],"recommendation":"major_revision","confidential_remarks":"The central idea is timely and the code release is a strength, but the paper's quantitative evidence needs substantial revision before it can be accepted. In particular, the unquantified exclusion of degenerate low-threshold configurations is a reporting issue that could amount to cherry-picking the Pareto frontier; this must be addressed head-on. The architectural comparison using Mistral vs Codestral is also a weakness because the models differ in pretraining domain, which the authors should either control or explicitly justify. I would encourage the editor to seek a revision where these points are fully addressed rather than reject, as the framework itself is plausible and the deficiencies are fixable within the manuscript's scope."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Quick take: this is the first paper I know to bolt early exits onto Mamba and to use a Mamba block as the exit predictor, and that is a legitimate, useful extension. But the headline empirical claim rests on a threshold-sweep plot where low-threshold configurations that produce degenerate repeated tokens are silently excluded, and the rule is unspecified. That is a load-bearing soft spot, not a nitpick.\n\nWhat's genuinely good: the framing is honest about prior work (CALM, EE-LLM), they release code, and the idea of using Mamba's recurrent state as a cheap confidence predictor is clever. The KV-cache copy trick for Transformers is a sensible small contribution, and the top-k relaxation of the oracle-matching objective is a reasonable training stabilizer. The savings are honestly reported as modest (~1.2x early-exit gains), which is consistent with prior EE work on decoder-only models.\n\nWhere it gets shaky: Section V says configurations are dropped when the model 'proceeds to output inconsistent responses that are composed of the same repeated token' and 'if this event occurs a certain number of times' the configuration is not shown. No count, no threshold, no sensitivity analysis. Since every plot's x-axis is the confidence threshold, removing the worst low-threshold points selects the Pareto frontier by construction. This directly inflates the savings-quality trade-off that the abstract claims. The rule applies only to EE configs, not layer pruning, so the comparison is asymmetric. I agree with the stress-test note: this is the load-bearing weakness.\n\nSecondary issues: Mistral 7B and Codestral Mamba 7B differ in pretraining data (Codestral is code-specialized), so architecture is confounded with pretraining. No error bars or repeated seeds anywhere. The KV-cache copy is described in a sentence but never specified: what exactly is copied, when, and does it affect attention masks? Also no wall-clock or energy measurements, only FLOPs-based reduction factors.\n\nBottom line: the novelty is real, the implementation is plausible, and the code is public, so a serious referee could pin these issues down. The paper deserves peer review, but the authors should be pushed to report excluded configurations, rerun with matched models or at least discuss the confound, add variance estimates, and specify the KV-copy mechanism. If the exclusion rule turns out to remove most low-threshold configs, the savings claim shrinks accordingly.\n\nI'd bring it to reading group, though it would be a discussion about evaluation standards as much as about the method.","headline":"First early-exit framework for Mamba with a clever Mamba-as-classifier design, but the headline savings curve rests on an unspecified rule that drops degenerate low-threshold runs.","tokens_in":11850,"tokens_out":2116,"would_cite":true,"duration_ms":21494,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Mamba's own blocks can cut LLM compute by deciding when to exit","keywords":["early exit","Mamba","state-space models","dynamic inference","decoder-only transformers","knowledge distillation","computational efficiency","large language models"],"falsifier":"Run DYNAMAX on a Mamba model with identical pretraining data and objective as a matched Transformer, or on the same Mamba weights before and after code specialization, and check whether the Mamba classifier's advantage over feed-forward exits persists; if the advantage vanishes, the claimed architectural benefit is confounded by pretraining.","tokens_in":10997,"feed_emoji":"⚡","tokens_out":3766,"duration_ms":36808,"temperature":0.7,"pith_summary":"DYNAMAX is a framework for adding early exits to decoder-only large language models, and the paper claims it is the first to do so for Mamba-based architectures. The core assertion is that Mamba blocks can serve not only as the backbone being exited but also as the confidence classifiers that decide when to exit, for both Mamba and Transformer hosts. If the claim holds, Mamba-based LLMs can stop computing per token as soon as confidence is sufficient, and the stateful structure of Mamba makes it a cheaper and more flexible exit predictor than feed-forward alternatives. The paper reports that early exits beat layer pruning on knowledge-focused QA tasks at similar compute savings, and that the Mamba classifier gives the widest range of useful confidence thresholds.","feed_headline":"Mamba's own blocks can cut LLM compute by deciding when to exit","feed_subtitle":"DYNAMAX attaches early-exit classifiers to Mamba and Transformer backbones, saving computation while keeping accuracy in QA tasks.","key_machinery":"The central object is the token-forwarding scheme of Algorithm 1 together with the three early-exit classifier heads: a CALM-style one-layer feed-forward net, a transformer-like FFN, and a Mamba block whose output projection is reduced to two logits. The load-bearing comparison is the operational-cost decomposition: a Transformer block costs about $24Td_{\\text{model}}^2 + 4T^2d_{\\text{model}}$ operations, while a Mamba block costs about $6d_{\\text{model}}^2 + 2n_{\\text{groups}}d_{\\text{state}}d_{\\text{model}}$, which the paper uses to argue that a Mamba classifier adds only constant per-token cost while carrying a recurrent state that helps it track the generated text. For Transformers, the efficiency trick is copying the KV cache entries from the exit layer directly to all subsequent layers instead of recomputing partial forwards for the exited token.","core_discovery":"The paper claims that Mamba-based large language models can be made dynamically computable: early-exit classifiers attached to the latter half of a Mamba backbone can stop per-token inference when confidence crosses a threshold, and, more strongly, that a Mamba block itself—with its output projection reduced to two logits—is a better early-exit classifier than a one-layer CALM-style head or a feed-forward block for both Mamba and Transformer hosts. The framework trains these exits by knowledge distillation against the full model's output, using a top-k relaxation of the oracle target to keep training stable, and reports that on TriviaQA, CoQA, and TruthfulQA the Mamba-classifier configuration preserves accuracy better than layer pruning at comparable compute savings, with the Mamba backbone showing a larger gap between recomputing and not recomputing states than the Transformer does.","pith_inferences":["If Mamba's recurrent state is what gives its exit classifier a wider threshold range, then hybrid or linear-attention models with explicit recurrent states might inherit the same benefit; this is an inference, not a paper claim.","The KV-cache copying trick could compose with speculative decoding or draft-model pipelines, where early-exited tokens are later re-verified with full computation.","A natural testable extension is per-layer adaptive thresholds driven by context or latency budgets, which the paper lists as future work but does not evaluate.","The top-k relaxation of the distillation target is a free parameter; varying k could expose a trade-off between training stability and exit accuracy that the paper does not quantify."],"forward_implications":["In Mamba-based LLMs, inference can be terminated per token when an exit confidence threshold is met, cutting compute on easy tokens and in low-latency settings.","A Mamba block with a two-logit output projection can act as a reusable early-exit classifier on Transformer backbones, not just on Mamba backbones.","For Transformer early exits, copying the KV cache from the exit layer to later layers avoids partial recomputation, lowering the overhead of exiting.","In knowledge-heavy tasks such as TriviaQA, early exits degrade accuracy more gracefully than layer pruning at the same compute reduction.","Skipping state recomputation in Mamba changes the cost-performance trade-off more than it does in Transformers, so the choice is architecture-specific."],"supporting_citations":[{"why":"Supplies the Mamba state-space architecture that the paper extends with early exits and repurposes as an exit classifier.","marker":"[11]"},{"why":"Provides the Mistral 7B Transformer baseline whose behavior is compared against the Mamba backbone.","marker":"[5]"},{"why":"Provides the Codestral Mamba 7B model, the Mamba-based backbone used in all experiments.","marker":"[38]"},{"why":"Defines the CALM early-exit framework whose style of classifier and oracle setting the paper adapts and compares against.","marker":"[24]"},{"why":"Supplies the layer-pruning baseline that early exits are compared with across all tasks.","marker":"[16]"},{"why":"Supplies the FineWeb-Edu corpus used to train the early-exit classifiers.","marker":"[40]"},{"why":"Supplies the evaluation harness and metrics used to score the models on TriviaQA, CoQA, and TruthfulQA.","marker":"[44]"}],"fun_headline_variants":["Mamba knows when to stop: early exits trim LLM compute","Let Mamba decide the exit point—save compute in LLMs","Mamba's own blocks can serve as cheap early-exit classifiers","DYNAMAX: dynamic exits for Mamba and Transformer LLMs","Two-logit Mamba heads make smart early-exit calls for LLMs"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The empirical case rests on comparing Mistral 7B with Codestral Mamba 7B as if they differ only in architecture, when the Mamba model was pretrained specifically for code tasks; differences in early-exit behavior could come from that pretraining gap.","fun_headline_variants_meta":{"raw":{"variants":["Mamba knows when to stop: early exits trim LLM compute","Let Mamba decide the exit point—save compute in LLMs","Mamba's own blocks can serve as cheap early-exit classifiers","DYNAMAX: dynamic exits for Mamba and Transformer LLMs","Two-logit Mamba heads make smart early-exit calls for LLMs"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000787,"raw_usage":{"total_tokens":3476,"prompt_tokens":955,"completion_tokens":2521,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":571,"completion_tokens_details":{"reasoning_tokens":2442}},"tokens_in":571,"tokens_out":2521,"duration_ms":19207,"temperature":1.0,"reasoning_tokens":2442,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-16T05:15:32.881865+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run DYNAMAX on a Mamba model with identical pretraining data and objective as a matched Transformer, or on the same Mamba weights before and after code specialization, and check whether the Mamba classifier's advantage over feed-forward exits persists; if the advantage vanishes, the claimed architectural benefit is confounded by pretraining.","supporting_citations":[{"cited_title":"Mamba: Linear-Time Sequence Modeling with Selective State Spaces,","cited_arxiv_id":null,"evidence_quote":"Supplies the Mamba state-space architecture that the paper extends with early exits and repurposes as an exit classifier."},{"cited_title":"Mistral 7B,","cited_arxiv_id":null,"evidence_quote":"Provides the Mistral 7B Transformer baseline whose behavior is compared against the Mamba backbone."},{"cited_title":"Codestral mamba,","cited_arxiv_id":null,"evidence_quote":"Provides the Codestral Mamba 7B model, the Mamba-based backbone used in all experiments."},{"cited_title":"The Unreasonable Ineffectiveness of the Deeper Layers,","cited_arxiv_id":null,"evidence_quote":"Supplies the layer-pruning baseline that early exits are compared with across all tasks."},{"cited_title":"Fineweb-edu,","cited_arxiv_id":null,"evidence_quote":"Supplies the FineWeb-Edu corpus used to train the early-exit classifiers."},{"cited_title":"A framework for few-shot language model evaluation,","cited_arxiv_id":null,"evidence_quote":"Supplies the evaluation harness and metrics used to score the models on TriviaQA, CoQA, and TruthfulQA."}],"review_version":1}