{"id":"8560b35a-2bd5-4511-89f2-0783ba85fa24","arxiv_id":"2502.09245","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":5.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"Adding learned per-head routing over previous layers' key-value buffers to a Transformer reduces representation collapse, lowers language modeling loss, and improves synthetic arithmetic and planning accuracy.","lead":"This paper proposes LIMe, a small addition to Transformer attention that lets each layer mix key-value memories from all earlier layers instead of only the previous layer's hidden state. This routed memory keeps token representations more distinct in deep layers, improving language modeling loss and giving large gains on synthetic arithmetic and planning tasks.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Collapse-mitigation claim is unsupported by the paper's own collapse metrics: LIMe hidden states are less separable than LLaMa's (Fig. 7b), and value-space proxies are never causally linked to the task gains.","rationale":"The reader's conditional verdict is appropriate, and my read supports it rather than moving it. The architecture is clearly described, the ablations in Appendix D show that the full router matters over fixed-average and windowed variants, and the comparisons against Hyper-Connections provide a relevant baseline. The concern is not that LIMe fails to improve loss or accuracy; the concern is that the paper's stated mechanism, mitigating representation collapse, is not supported by the evidence. The hidden-state separability result goes in the opposite direction from the cited collapse definition, and the value-space metrics are never shown to be causal. This is a correctness risk for the interpretation, not for the raw empirical claims. A single intervention, matching value entropy in the baseline, would settle whether the proxy is causal, which is why I would keep the paper conditional rather than reject it.","tokens_in":20082,"tokens_out":12847,"duration_ms":130589,"concrete_test":"Run a controlled intervention on the Arithmetic Expression Task: train the LLaMa 4-layer baseline with a value-entropy regularizer (e.g., Seq-VCR's variance-covariance penalty) so that its value-vector Renyi entropy matches LIMe's, while keeping the architecture and all other training conditions fixed. If accuracy does not rise toward LIMe's 71.6% on 6 operands, the value-entropy proxy is not the causal channel behind LIMe's gains and the collapse-mitigation explanation fails.","verdict_should_be":"UNCHANGED","load_bearing_attack":"Section 5.2's collapse evidence is internally two-sided and does not establish the mechanism stated in the abstract. The paper cites Voita et al. (2019) and Barbero et al. (2024) for collapse as a loss of fine-grained token distinctions in deeper-layer representations; by that standard metric, LIMe is worse, not better: Figure 7(b) shows lower linear separability of hidden states for LIMe than for LLaMa, and Figure 7(a) shows no entropy difference in hidden states. The paper responds that hidden states need not pack information because values carry it, but this is a redefinition rather than a measurement of collapse. The only favorable metrics are value-vector entropy and value separability (Figs. 2a, 2b), and no experiment connects those proxies to next-token loss or arithmetic accuracy. Because Eq. (2) ties keys and values to the same learned router weight, value-space diversity is not an independently controlled variable; it could be a side effect of linearly mixing many layers' K/V buffers. The central explanatory claim therefore rests on an untested correlation, even though the empirical gains themselves may be real.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes Layer-Integrated Memory (LIMe), a modification of decoder-only Transformers in which each attention head at each layer learns a per-head, per-layer routing weight that mixes the key and value buffers of all previous layers (Eq. 2), instead of using only the current layer's projected K/V. The authors claim that standard Transformers suffer from representation collapse because they rely on the previous layer's hidden state, and that LIMe mitigates this collapse, leading to faster convergence per FLOP, lower perplexity, higher accuracy on ProsQA and arithmetic reasoning, and better value-space entropy/separability. Experiments include 1B-parameter language modeling, fine-tuning on ProsQA, training on arithmetic expressions, deep 32/64/128-layer training-loss comparisons, and a router ablation.","tokens_in":20328,"tokens_out":10779,"duration_ms":98384,"significance":"If validated, LIMe is a simple and plausible architectural fix: it reuses existing key-value buffers, adds negligible parameter overhead, and comes with clear pseudocode, a public code repository, and a broad set of experiments including GQA/full attention settings and comparisons with Hyper-Connections. The cross-layer routing idea is well motivated and the empirical gains on synthetic reasoning are potentially meaningful. However, the paper's central explanatory claim--that LIMe mitigates representation collapse--is not established by the evidence, and the reported efficiency numbers are internally inconsistent with the described O(L^2) router. The deep-network result (64-layer LIMe beats 128-layer LLaMA) rests only on training loss. These issues currently make the main claims uncertain, but they are addressable with additional analysis and clarifications.","major_comments":[{"comment":"The collapse-mitigation claim is not supported by the paper's own metrics. Figure 7(b) shows that LIMe hidden states are less linearly separable than LLaMA's, which by the definition of collapse cited in Section 2 (Voita et al., 2019: deeper layers lose fine-grained token distinctions) indicates more collapse, not less. The paper responds that hidden states need not carry all information because values do, but this is a re-interpretation rather than a measurement of the original phenomenon. Moreover, no experiment connects the favorable metrics (value entropy in Fig. 2(a), value separability in Fig. 2(b)) to next-token loss or synthetic-task accuracy. Because Eq. (2) uses the same learned router weights for keys and values, the reported value-space diversity could be a byproduct of linearly mixing many layers' buffers rather than the causal mechanism behind the task gains. To support the mechanistic claim, the authors should test the link directly, e.g., by intervening on the router to reduce value entropy/separability and measuring the effect on downstream accuracy, or by evaluating final-layer hidden-state separability for the specific function words used in the classifier.","section":"Section 5.2, Figures 2, 3, 7"},{"comment":"The deep-network claim that 'a 64-layer LIMe model outperforms a 128-layer LLaMA model' is load-bearing for the scaling conclusion, but it rests entirely on training-loss curves at 10B tokens, with no held-out validation loss or downstream evaluation. The section does not report the hidden dimension, KV-head configuration, training data, or total compute for these runs, and there are no multiple seeds or error bars. Please add evaluation on a held-out set, specify the exact architecture and compute budget, and report run-to-run variation; otherwise the claim that LIMe changes scaling behavior is not substantiated.","section":"Section 5.5, Figure 6"},{"comment":"The reported FLOP overhead is inconsistent with the routing operation described in Eq. (2) and the pseudocode in Appendix E. In the pseudocode, the router at layer l performs a dense matmul between a [kv_h, (l+1)*kv_h] weight matrix and the full KV buffer, which costs O(l) per layer and O(L^2) overall. For the stated dimensions (L=16, d=2048, H=32, kv_h=8, dh=64) and a training sequence of 2048 tokens, this router matmul accounts for a much larger share of total FLOPs than the +0.08% (GQA) and +1.22% (full attention) reported in Table 8. This matters because the central iso-FLOP comparison (Fig. 1) and the 'negligible overhead' claim depend on how router FLOPs are counted. Please clarify whether the reported FLOPs are per-step or per-token, how the router matmul was included (or excluded), and provide wall-clock step times for the actual training configuration.","section":"Table 8, Eq. (2), Appendix E"},{"comment":"Table 7 reports that the 'LIMe average' variant achieves validation perplexity 16.4611, exactly identical to the LLaMA baseline (16.4611). Since a uniform average of all buffered key-value representations necessarily changes the attention inputs relative to a standard transformer, an exactly equal perplexity is implausible and suggests a bug in the experiment or in the reporting. This makes the ablation's conclusion--that uniform pooling of past representations is insufficient--uninterpretable. Please correct the table or explain why the two rows coincide.","section":"Table 7, Appendix D"},{"comment":"All main empirical comparisons are based on single runs, with no confidence intervals, multiple seeds, or significance tests. The claimed margins are small (e.g., 1.15% perplexity difference in Section 5.1, 3.36% in Table 7, 8.4% on ProsQA) and the paper describes the results as 'consistently achieved'. Without variance information, the reader cannot judge whether these differences are within run-to-run noise. Please provide at least 3-5 seeds for the primary comparisons (language modeling loss, ProsQA, arithmetic, and the deep-network training curves), or state explicitly whether the training is deterministic and why.","section":"Tables 1, 4-6; Figures 1, 6; Section 5.3"}],"minor_comments":[{"comment":"There are repeated typographical glitches in author names and math: 'V oita' appears in the abstract-related text and in the references, and 'R ´enyi' has a spurious accent. Please fix these.","section":"Throughout"},{"comment":"The notation using H for the number of heads is later reused as the second and third dimensions of the router tensor R^(l) in Eq. (2), but in the GQA experiments the router mixes only num_kv_heads (8) rather than num_heads (32), as shown in Appendix E. Define whether Eq. (2)'s sums over h' run over all heads or over KV heads, and make the pseudocode consistent with the equations.","section":"Section 3, Eq. (2)"},{"comment":"The precise definitions of '15.3% (8.9% with GQA) faster in FLOPs' and '1.15% (0.91% with GQA) lower perplexity' are not given. Please state the reference loss value and how the convergence speed is computed (e.g., FLOPs to reach a fixed loss).","section":"Section 5.1"},{"comment":"The sentence 'Unlike LLaMA, LIMe can make updates attending to the previous representations' is unclear; rephrase to say what the t-SNE panels show about hidden-state versus value-state separability.","section":"Figure 3 caption"},{"comment":"The router learning rate is set to 0.01 while the base LR is 0.001, and the identity initialization is described as important (Appendix A). No sensitivity analysis is reported for these hyperparameters; a brief ablation or a comment on stability would help.","section":"Table 2, Appendix A"},{"comment":"The analysis of the router weights uses absolute magnitudes because the sign is 'semantically ambiguous'. While reasonable, this discards the distinction between positive and negative routing, which may matter for interpretation. Please add at least one qualitative check on the signed weights or a small case study.","section":"Section 5.4, Figure 5"}],"recommendation":"major_revision","confidential_remarks":"The paper has a clear architecture and a plausible set of experiments, but several load-bearing issues remain. The FLOP accounting in Table 8 appears difficult to reconcile with the O(L^2) router implementation described in Eq. (2) and Appendix E; if the FLOP curve in Fig. 1 was computed without including the router matmuls, the 'converges faster per FLOP' claim would be overstated. The exact equality of the 'LIMe average' and LLaMA perplexity rows in Table 7 is a red flag for the ablation reliability. The collapse-mechanism framing also oversells Fig. 7(b), where LIMe's hidden states are less separable; I would urge the editor to ask for the intervention-style causal evidence suggested in my major comments rather than accepting the current correlation-based narrative. If these points are resolved in revision, the work could be a solid contribution."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"LIMe is a cheap cross-layer routing mechanism that deserves a serious look, but the paper's central collapse narrative is the weakest part. The method is a genuine incremental twist: instead of attending over previous hidden states, it learns per-head, per-layer weights over the existing KV buffers, so the extra FLOP and memory costs are tiny (under 1% in GQA, by their numbers). The experiment suite is broad—1B models trained on 50B tokens, LM Eval, ProsQA, arithmetic expressions, and a 32/64/128-layer scaling study—and the synthetic arithmetic gain (+30% at 6 operands) is notable if it holds. The router analysis showing systematic reuse of early and neighboring layers is a useful sanity check.\n\nThe soft spots are real and concentrated where the paper makes its strongest claim. The representation-collapse evidence is internally two-sided: hidden states in LIMe are less separable than LLaMa's (Fig. 7b), and the authors' response—'values carry the information'—is a redefinition, not a measurement of the collapse they say they mitigate. No experiment links value-vector entropy or separability to next-token loss or arithmetic accuracy, and because Eq. (2) shares one learned weight between keys and values, value diversity is not an independent knob. The deep-network claim ('64-layer LIMe matches 128-layer LLaMa') is based only on training loss; the main benchmark tables are single runs without confidence intervals or seeds; and Table 7 contains an exact duplicate perplexity value between LLaMA and the LIMe-average variant that looks like a copy-paste error. The O(L^2) complexity and pipeline communication overhead are acknowledged but could matter at scale.\n\nI would not reject the paper. The architecture is plausible and the empirical gains, while underpowered, are consistent across several settings. But the abstract overstates the mechanism. A referee should ask for multi-seed runs, clarification of the duplicate, and either a causal experiment (e.g., freezing the router or artificially controlling value diversity) or a toned-down collapse claim. The paper is worth engaging with; it just needs revision before I would trust the numbers.","headline":"Useful cross-layer routing with a plausible but unproven collapse story; send to review with demands for seeds and causal tests.","tokens_in":20853,"tokens_out":2715,"would_cite":false,"duration_ms":26137,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Standard transformers collapse context into one hidden state, blurring distinctions at depth; LIMe routes attention through all earlier key-value buffers, cutting convergence cost 15.3% and lifting reasoning by 30 points.","keywords":["representation collapse","layer-integrated memory","inter-layer routing","key-value buffers","transformer architecture","multi-step reasoning","Rényi entropy","depth scaling"],"falsifier":"Train a LLaMa baseline with an auxiliary objective that drives its value-vector Rényi entropy and token separability up to LIMe's measured levels; if its accuracy on ProsQA and six-operand arithmetic does not rise correspondingly, then the entropy/collapse channel is not the causal path behind LIMe's gains, and some other feature of the routing — extra gradient pathways, initialization, or the shared key-value weights — would be responsible.","tokens_in":19884,"feed_emoji":"🧠","tokens_out":19239,"duration_ms":148392,"temperature":0.7,"pith_summary":"Standard transformer decoders compress the whole context into the hidden state produced by the immediately previous layer. This paper argues that this single-stream design is a genuine flaw: in deep networks it induces representation collapse, with fine distinctions between tokens blurring until they become hard to separate, and the cost shows up as worse loss and failed multi-step reasoning. The proposed fix, Layer-Integrated Memory (LIMe), reuses the key-value buffers a transformer already keeps and learns per-head, per-layer routing weights so that each attention head can mix representations from all earlier layers, with essentially no added memory and around a percent of extra FLOPs. The paper reports that a 1B-parameter LIMe model converges 15.3% faster per FLOP, reaches 1.15% lower perplexity, gains 8.4 points on the ProsQA search task and 30 points on six-operand arithmetic, and makes a 64-layer model match a 128-layer baseline. The aim is to show that representation collapse, not depth or width alone, is the bottleneck in deep transformers, and that learned inter-layer memory removes much of it almost for free.","feed_headline":"Route attention through all earlier layers to beat 2x-deeper nets","feed_subtitle":"Reusing existing key-value buffers cuts convergence cost 15.3 percent and lifts multi-step math accuracy by 30 points.","key_machinery":"The load-bearing object is the per-layer, per-head router tensor $R^{(\\ell)} \\in \\mathbb{R}^{\\ell \\times H \\times H}$, whose entry $R^{(\\ell)}_{\\ell',h',h}$ weights the contribution of head $h'$ at buffered layer $\\ell'$ into head $h$ at layer $\\ell$. It forms routed keys and values by the same weighted sum, $\\tilde{K}_{\\ell,h} = \\sum_{\\ell'=1}^{\\ell}\\sum_{h'=1}^{H} R^{(\\ell)}_{\\ell',h',h} B^{(K)}_{\\ell',h'}$ and an analogous expression for values, then runs ordinary causal attention between the current layer's queries and that routed mixture. Because the buffers $B^{(K)}$ and $B^{(V)}$ are the ones vanilla transformers already maintain, the mechanism adds almost no memory, and because the router is a single matrix multiply, the added compute stays a small fraction. Initializing the diagonal slice to the identity and the remaining entries to Kaiming-uniform values lets the network start as a standard transformer and gradually open cross-layer pathways; the appendix ablation shows that replacing the learned router with a uniform average over past layers reproduces the baseline perplexity exactly, so the learned weighting is what carries the benefit.","core_discovery":"The central claim is that a standard decoder's habit of feeding each attention layer only the previous layer's hidden state is not a harmless design detail but the source of representation collapse: as layers deepen, the residual stream must carry every contextual distinction forward in a single vector, and under finite precision and capacity those distinctions merge, so similar tokens become indistinguishable exactly where the model needs them separated. LIMe's discovery is that the key-value buffers a transformer already allocates can serve as a persistent multi-layer memory: a small trainable router tensor per layer and head, initialized to the identity, learns to weight the buffered key-value states of all earlier layers and feeds the mixture into ordinary attention. The empirical case is that this changes measurable behavior — higher value-vector Rényi entropy, near-perfect linear separability of function words, and better separation of nearby numbers in arithmetic — and that this translates into faster convergence, lower perplexity per FLOP, large gains on ProsQA and the Arithmetic Expression Task, and depth scaling in which a 64-layer LIMe outperforms a 128-layer baseline. Analysis of the trained routers shows systematic reuse of embedding-layer features in early layers and retrieval of the first buffers in later layers, which the authors read as evidence that the inter-layer pathways, rather than extra parameters alone, are doing the work.","pith_inferences":["A test the paper leaves open is whether keys and values should share one router: Equation (2) couples them with a single tensor, and decoupling the two streams could reveal whether the value stream alone drives the gains or whether independent key routing helps further.","The routing heatmap hints that a sparse router — attending mainly to the first few buffers plus the immediate predecessor — might capture most of the benefit; this matters because the paper's limitations section concedes the vanilla router has $O(L^2)$ cost and adds GPU communication in pipeline-parallel training, both pressing against the depth scaling the method promises.","The entropy and separability diagnostics may be symptoms of routing rather than its causal channel; an intervention that raises a baseline's value-vector entropy without adding routing would decide whether 'mitigating collapse' explains the gains or merely accompanies them.","If the mechanism holds, LIMe's less-collapsed value states are a natural substrate for continuous latent-space reasoning, letting a model 'think' in earlier layers that have not yet squashed distinctions — a direction the paper flags for future work but does not test."],"forward_implications":["Depth scaling changes: a 64-layer LIMe model reaches lower training loss than a 128-layer LLaMa baseline, so the same compute budget buys more capability when deep layers can read earlier memory instead of carrying everything in the residual stream.","Multi-step reasoning improves on open-ended generation: +8.4 points on ProsQA (77.8% versus 69.4%) and about 30 points on six-operand arithmetic (71.6% versus 41.3%), where the baseline's intermediate values collapse.","The added cost is minimal: roughly 0.01% more parameters, 0.08% more FLOPs with grouped-query attention (1.22% without), essentially unchanged peak memory, and compatibility with efficient attention implementations such as FlashAttention.","The trained routers are used as designed: early layers lean on the embedding buffers, middle layers treat the immediate predecessor as an auxiliary memory bank, and final layers reach back to the first buffers, so the extra pathways are actually exercised rather than idle."],"supporting_citations":[{"why":"Supplies the Rényi-entropy collapse metric, the variance-covariance regularization framing, and the Arithmetic Expression Task behind the 30-point reasoning result.","marker":"[Arefin et al., 2024]"},{"why":"Supplies the four-way linear-separability test on is/are/was/were that the paper uses to show collapse is avoided in LIMe's value vectors.","marker":"[Voita et al., 2019]"},{"why":"Gives the theoretical information over-squashing result for decoder-only transformers that motivates why deep layers collapse representations.","marker":"[Barbero et al., 2024]"},{"why":"Defines the LLaMa architecture family that LIMe modifies and that serves as the main baseline for loss, benchmarks, and efficiency.","marker":"[Grattafiori et al., 2024]"},{"why":"Supplies ProsQA, the graph-search reasoning benchmark on which LIMe gains 8.4 points over the baseline.","marker":"[Hao et al., 2024]"},{"why":"Co-introduces the arithmetic expression solution-step task used, with Arefin et al., to test multi-step intermediate computation.","marker":"[Feng et al., 2023]"},{"why":"Hyper-Connections is the competing collapse-mitigation baseline; its cost and accuracy are the benchmark for LIMe's negligible-overhead claim.","marker":"[Zhu et al., 2024]"}],"fun_headline_variants":["Reuse KV buffers across layers: beat 2x-deeper transformers","Stop feeding layers only one hidden state; reuse all keys and values","Transformer capacity wasted? Route attention across all prior layers","LIMe: reuse previous layers' key-values for richer representations","Add per-layer routing to key-value caches; beat 128-layer baselines"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that the collapse measurements — the Rényi entropy of value vectors and the linear separability of four verb forms — capture the mechanism behind LIMe's task gains; the paper shows they move together with performance but never shows that raising entropy or separability is what lowers loss or raises accuracy.","fun_headline_variants_meta":{"raw":{"variants":["Reuse KV buffers across layers: beat 2x-deeper transformers","Stop feeding layers only one hidden state; reuse all keys and values","Transformer capacity wasted? Route attention across all prior layers","LIMe: reuse previous layers' key-values for richer representations","Add per-layer routing to key-value caches; beat 128-layer baselines"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.001312,"raw_usage":{"total_tokens":5366,"prompt_tokens":983,"completion_tokens":4383,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":599,"completion_tokens_details":{"reasoning_tokens":4292}},"tokens_in":599,"tokens_out":4383,"duration_ms":27665,"temperature":1.0,"reasoning_tokens":4292,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-07T22:11:08.838610+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Train a LLaMa baseline with an auxiliary objective that drives its value-vector Rényi entropy and token separability up to LIMe's measured levels; if its accuracy on ProsQA and six-operand arithmetic does not rise correspondingly, then the entropy/collapse channel is not the causal path behind LIMe's gains, and some other feature of the routing — extra gradient pathways, initialization, or the shared key-value weights — would be responsible.","supporting_citations":[],"review_version":1}