{"id":"44b6e8e9-e439-4b41-bbb5-f84fe2aaa4d2","arxiv_id":"2506.05229","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":0,"one_line_summary":"Diagonal Batching regroups layer-segment computations into parallel diagonals, preserving exact recurrence and speeding up long-context recurrent memory inference without retraining.","lead":"This paper introduces a scheduling scheme for recurrent memory transformers that processes several segments at once by grouping the segment-by-layer computation grid into diagonals. It reports up to 3.9x faster inference than the sequential recurrent implementation and up to 3.3x faster than full-attention Llama at 131,072 tokens, with no retraining.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The scheduling lemma is correct for PRMT/ARMT, but the paper's central claim is scoped to a layer-level DAG that standard RMTs, named in the title and abstract, do not satisfy; the title and the phrase 'existing RMT models adopt it' overstate the method's applicability.","rationale":"The reader's weakest assumption correctly identifies the scope limitation: Diagonal Batching relies on a layer-level recurrent dependency graph that standard RMTs violate. The scheduling lemma is elementary and correct for the PRMT/ARMT family, and the empirical speedups are plausible and consistently reported across model sizes. The paper is transparent about the restriction in Section 2.2 and Limitations, so this is not an internal inconsistency in the technical argument for PRMTs. However, the title and abstract present the method as applicable to RMTs generally and state that existing RMT models adopt it with no retraining, which is not true for the original RMT formulation. The concrete dependency-graph check would definitively show why the diagonal decomposition is invalid for standard RMTs. The reader's other conditions, such as clarifying the numerical drift in Table 2 and pinning the code, remain appropriate; the Appendix's note about a code discrepancy behind the drift is an additional reason to require the drift bound before the 'exact recurrence' wording is accepted. Overall, the central contribution for PRMTs/ARMT stands, so the conditional verdict should be unchanged.","tokens_in":13534,"tokens_out":23480,"duration_ms":220455,"concrete_test":"Construct the true dependency graph for the standard RMT with L=2 layers and S=4 segments using Equation (1): for each node (s, l), include every dependency implied by feeding the previous segment's final-layer memory into all layers of the next segment. Then compute the earliest feasible group for node (2, 1). Under the true RMT graph, this node depends on the final-layer memory output of segment 1, i.e. node (1, 2), so its earliest group is at least 3; Lemma 3.1 would instead assign it to group 2. Verifying this dependency graph directly settles whether the diagonal schedule preserves exact recurrence for standard RMTs.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The load-bearing assumption is the layer-level dependency graph in Section 3.1: each (segment, layer) node depends only on (segment-1, layer) and (segment, layer-1). This is exactly the PRMT/ARMT structure described in Section 2.2, but it is not the dependency graph of the original RMT defined by Equation (1) and Figure 2 (left), where the final-layer memory of the previous segment is fed into all layers of the next segment. For standard RMT, node (s, l) transitively depends on (s-1, L), the final layer of the previous segment, so its earliest feasible group is not s+l but at least s-1+L. Lemma 3.1, and the entire diagonal schedule, therefore does not preserve exact recurrence for the model class named in the paper's title. The paper explicitly discloses this restriction in Section 2.2 and in Limitations, which is a credit to its honesty, but the abstract's opening 'in RMTs' and the claim that 'existing RMT models adopt it with no retraining' are unsupported for standard RMTs. This makes the central claim, as advertised, too broad; the technical contribution is valid only for PRMTs/ARMT.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes Diagonal Batching, a scheduling scheme for layer-level recurrent models (which the authors call PRMTs, including ARMT, RWKV, and Mamba) that reorganizes the layer-segment computation grid into diagonal groups. The authors prove in Lemma 3.1 that this schedule completes the dependency DAG in the minimum number of groups, N_segments + N_layers - 1, and implement it for ARMT by replacing all layers with a single grouped layer. They report speedups over sequential ARMT and over standard full-attention Llama models at sequence lengths up to 131,072 tokens on A100/H100 GPUs, together with an error-accumulation study and BABILong evaluations.","tokens_in":13793,"tokens_out":7880,"duration_ms":80677,"significance":"If the claims hold, this is a useful systems contribution: it shows that a simple, purely runtime reordering can unlock inter-segment parallelism for PRMT-style architectures without retraining, and it ships publicly available code. The scheduling lemma is correct for the stated layer-level dependency graph, and the experimental sweep across model sizes (160M to 8B), sequence lengths, and GPUs is broad. The measurements are accompanied by an error analysis and a downstream benchmark, which strengthens the empirical picture. The main caveat is that the method does not apply to the original RMT recurrence, so the paper's broad 'RMT' framing overstates the contribution, and the 'exact recurrence' claim needs qualification given the reported numerical drift.","major_comments":[{"comment":"The title and abstract claim applicability to 'Recurrent Memory Transformers' generally, but Section 2.2 and the Limitations explicitly state that standard RMT, defined by Eq. (1), does not satisfy the layer-level dependency assumption on which Diagonal Batching relies. In particular, the abstract's statement that 'existing RMT models adopt it with no retraining' is unsupported for the original RMT, because node (s, l) in standard RMT transitively depends on the final layer of the previous segment, not only on (s-1, l) and (s, l-1). This is a load-bearing overstatement of the method's scope; the paper should consistently frame the contribution as applying to layer-level recurrent models (PRMTs/ARMT) in the title, abstract, and conclusion, or should justify why the original RMT can be included despite the dependency structure shown in Figure 2 (left).","section":"Title, Abstract, and Section 2.2"},{"comment":"The paper repeatedly claims that Diagonal Batching 'preserves exact recurrence' and enables 'exact, linear-time inference,' but Table 2 reports up to 1.87% relative logit drift between the base ARMT implementation and the Diagonal Batching implementation, and Appendix A attributes this drift to implementation details such as optimized kernels. The schedule does preserve the dependency graph exactly, but the executed computation is not bitwise identical. The authors should explicitly distinguish between exact preservation of the recurrence structure and numerical equality of the logits, and adjust the wording in the abstract and introduction accordingly; otherwise the 'exact' claim is misleading.","section":"Abstract, Section 3.2, and Table 2"},{"comment":"The Llama-3.2-1B full-attention baseline timing at 8192 tokens (0.026s) is internally inconsistent: it is barely larger than the 4096-token time (0.024s) and far too small relative to the 16384-token time (0.376s), which is close to the ~4x value expected under quadratic scaling from 8192 at ~0.096s. This inconsistency affects the credibility of the baseline and of the speedup figures derived from it, such as those in Table 8. The authors should correct the entry or explain the measurement conditions (e.g., a typo or a caching effect), and ideally report the number of trials and variance for the timing measurements.","section":"Table 1"}],"minor_comments":[{"comment":"The sentence 'However, we the effect of error accumulation on downstream tasks is negligible' appears to be missing a verb; it should read something like 'However, we find that the effect of error accumulation on downstream tasks is negligible.'","section":"Section 4.5"},{"comment":"The notation in Eq. (1) is confusing: the input list [M_{s-1}, H_{s-1}, M_{s-1}] repeats the memory state, and the output is written as [_, _, M_s]. This looks like a typesetting artifact and should be corrected to match the original RMT formulation or clarified in the text.","section":"Section 2.1, Eq. (1)"},{"comment":"The 'Ideal Even Load' curve is not precisely defined in the main text; the caption says it assumes all segment computations run with maximum achievable FLOPS, but the authors should specify how this ideal time was computed (e.g., measured peak FLOPS vs. Achieved FLOPS) so that the comparison is reproducible.","section":"Figure 6"},{"comment":"The table captions say 'speedup' but the columns contain values that appear to be execution times (e.g., 0.085s at 4096 tokens in Table 8); please clarify whether the tables report absolute times or speedup ratios, and label the columns consistently.","section":"Tables 8 and 9"},{"comment":"There are several minor typos and inconsistent capitalizations across the paper, including 'LLama' versus 'Llama' and 'inplementation' in the Table 9 caption; a careful proofreading pass is recommended.","section":"Global"}],"recommendation":"major_revision","confidential_remarks":"The paper's core scheduling idea is correct for PRMTs/ARMT and the experiments are extensive, but the overbroad 'RMT' framing and the 'exact' wording will likely draw criticism from readers familiar with the original RMT architecture. The authors can address this with a focused revision that re-scopes the title/abstract and clarifies the exactness claim; the Table 1 inconsistency also needs fixing before acceptance. I see no circularity concern: the method is an algorithmic/empirical contribution with no fitted constants, and self-citation of [28] is an architectural premise, not a circular dependence."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Thanks for the report. I read the paper. My take: the core scheduling idea is correct and the implementation is real, but the package has a scoping mismatch and a couple of presentation issues.\n\nWhat's new: applying anti-diagonal (wavefront) scheduling to the (segment, layer) dependency graph of PRMT/ARMT inference, then batching the operations in each diagonal with grouped GEMM and batched attention. Lemma 3.1 is elementary—topological sort, longest path—but it's the right framework, and it's new in this context as far as I know. The speedups are plausible: up to 2.7x over sequential ARMT for 1B at 128k, up to 3.9x for the 160M model, and they hold across model sizes and segment sizes. The BABILong parity check is a good touch; it shows the numerical drift doesn't hurt task performance. The code being public is also a plus.\n\nThe soft spots. First, the title and abstract say 'Recurrent Memory Transformers' but the method only applies to layer-level recurrent models (PRMT/ARMT), not the original RMT of Bulatov et al., where final-layer memory feeds all layers of the next segment. The body and Limitations are explicit about this, but the abstract's 'existing RMT models adopt it with no retraining' is too broad. That should be fixed before publication. Second, Table 1's full-attention Llama-1B row has a suspicious jump: 0.026s at 8k to 0.376s at 16k. That looks like a typo (maybe 0.076?). The authors should correct it. Third, the 'exact recurrence' language sits uneasily with the 1.87% logit drift shown in Table 2. The scheduling is mathematically exact; the drift is numerical, from grouped GEMM and different summation order. They do compare it to FlashAttention drift and show quality parity, so it's a wording issue, but the abstract's 'exact' and conclusion's '1% error' in the same breath is sloppy.\n\nBottom line: this is a solid engineering contribution, not a breakthrough. The scheduling lemma is correct, the measurements are broadly consistent, and the limitations are honestly disclosed. I'd send it to a main venue as a systems/ML paper; a referee should ask for the scoping fix, the Table 1 correction, and a more careful use of 'exact'. The theory is too thin for a theory venue, but it doesn't need to be.\n\nYes, I'd cite it if I worked on recurrent memory inference. I'd recommend peer review with minor revisions.","headline":"Solid engineering contribution with a correct scheduling lemma, honestly scoped in the limitations but over-sold in the title and abstract; worth a serious referee after some fixes.","tokens_in":14326,"tokens_out":3530,"would_cite":true,"duration_ms":40354,"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":"Diagonal Batching reorders layer-level recurrent memory transformer computation into diagonal waves, running up to $N_{\\text{layers}}$ operations in one GPU kernel launch and cutting 131,072-token inference latency by 3.3x.","keywords":["Diagonal Batching","Recurrent Memory Transformer","Parallel Recurrent Memory Transformer","Associative Recurrent Memory Transformer","long-context inference","GPU scheduling","segment-level recurrence","linear-time inference"],"falsifier":"Run one 131,072-token prompt through the same trained ARMT checkpoint twice, once with the sequential reference implementation and once with Diagonal Batching, on the same GPU, and compare final logits and wall-clock latency. If the relative logit error grows well beyond the reported roughly 2% as the number of segments increases, or if the grouped schedule is not faster than the sequential one for the claimed configurations, the exact-recurrence and speedup claims would be refuted. A direct boundary check is to apply the schedule to an original RMT with cross-layer memory flow and observe the output change, since the method is only valid for layer-level recurrent models.","tokens_in":13339,"feed_emoji":"⚡","tokens_out":13327,"duration_ms":138615,"temperature":0.7,"pith_summary":"The paper's claim is that the long-context inference bottleneck in Recurrent Memory Transformer–style models is scheduling, not algorithmic complexity: recurrent dependencies force layers and segments to wait for one another and leave GPUs under-used. Diagonal Batching is a run-time reordering of the two-dimensional grid of layers and segments into diagonal waves, so all independent computations with the same segment-plus-layer index run together; up to $N_{\\text{layers}}$ operations execute per GPU kernel launch while the exact recurrence is preserved. Because it is only a reordering, existing Parallel Recurrent Memory Transformers (PRMTs), such as the Associative Recurrent Memory Transformer (ARMT), adopt it with no retraining. On a one-billion-parameter ARMT, the paper reports a 3.3x speedup over full-attention inference and a 1.8x speedup over the sequential ARMT implementation at 131,072 tokens, with numerical drift under 2% and unchanged scores on a long-context reasoning benchmark.","feed_headline":"Diagonal batching runs recurrent-memory LLMs 3.3x faster","feed_subtitle":"It groups independent layer-segment computations into diagonals, preserving exact recurrence with no retraining.","key_machinery":"The central object is Diagonal Batching, a schedule for the layer-by-segment dependency graph. In a PRMT, node $(s, l)$ (segment $s$, layer $l$) depends only on $(s, l-1)$ and $(s-1, l)$, so all nodes with $s+l$ equal to the same value are mutually independent; Diagonal Batching groups them into one wave and launches the whole wave through a single grouped layer (stacked weights, grouped matrix multiplication, and batched attention). This makes the GPU treat the diagonal as a batch, recovering the utilization that normally requires many independent requests, without changing the recurrence the model computes.","core_discovery":"The central claim is that for models with layer-level recurrence—where each layer keeps its own memory and updates it once per segment—the whole forward pass over $N_{\\text{segments}}$ segments and $N_{\\text{layers}}$ layers can be executed as $N_{\\text{segments}}+N_{\\text{layers}}-1$ groups instead of $N_{\\text{segments}}\\times N_{\\text{layers}}$ sequential steps. Group $i$ contains every node whose segment index plus layer index equals $i$, and Lemma 3.1 states that this is the minimum possible number of groups because the longest dependency path has exactly that many nodes. The implementation fuses the distinct transformer layers into one grouped layer with stacked weights and batched attention, so the GPU sees a large effective batch without batching multiple requests. The paper argues that the resulting computation is numerically close to the original—relative logit error below 2% for sequences up to 32,768 tokens—and that on a 1B ARMT it delivers up to 3.3x lower latency than full-attention inference and up to 1.8x lower latency than the sequential ARMT at 131,072 tokens.","pith_inferences":["Editorial extension: the diagonal schedule is a generic property of any two-dimensional layer-by-segment dependency grid, so it should transfer to chunked inference in other per-layer recurrent architectures whenever their layer states update independently per segment.","Editorial extension: since the reported speedup grows with layer count, stacking more layers or designing grouped modular layers could compound the gain, while very shallow models will see little benefit.","Editorial extension: the backward pass described in the appendix suggests the same grouped schedule could be used during training, making long-context recurrent-memory training numerically consistent with inference and possibly faster.","Editorial extension: the observed 1–2% error budget is similar to replacing one attention kernel with another, so combining diagonal batching with quantization or speculative decoding is a plausible next test, though the paper does not demonstrate that combination."],"forward_implications":["Any already-trained PRMT-style model can be switched to Diagonal Batching as a drop-in inference change, with no retraining or architecture modification, because the reordering computes the same recurrence.","The number of sequential synchronization barriers drops from $N_{\\text{segments}} \\times N_{\\text{layers}}$ to $N_{\\text{segments}} + N_{\\text{layers}} - 1$, so models with more layers gain disproportionately more.","A single long-context request can saturate the GPU without waiting for other requests, which simplifies serving and load balancing compared with large-batch inference.","Developers can choose smaller segment sizes for quality without automatically paying the usual inference-speed penalty, because diagonal grouping decouples performance from segment size.","The observed numerical drift stays below 2% and generation scores on a long-context reasoning benchmark are unchanged up to 64k tokens, so the speedup is not bought with task quality."],"supporting_citations":[{"why":"Defines the Associative Recurrent Memory Transformer and the broader parallel-recurrent family whose layer-segment dependency graph the method schedules.","marker":"[28]"},{"why":"Introduces the original Recurrent Memory Transformer with cross-layer memory flow, the contrast case diagonal batching cannot accelerate.","marker":"[3]"},{"why":"Supplies the optimized batched attention kernel used inside the grouped forward pass and the 1–2% numerical-error reference level.","marker":"[5]"},{"why":"Provides the long-context reasoning benchmark used to show generation quality is preserved when the reordering is applied.","marker":"[17]"},{"why":"Defines the transformer model family used as the base architecture for the ARMT experiments.","marker":"[11]"},{"why":"Cited as an example of a layer-level recurrent linear architecture that fits the parallel-recurrent assumption.","marker":"[12]"},{"why":"Cited as another layer-level recurrent architecture whose per-layer memory would fit the diagonal schedule.","marker":"[24]"}],"fun_headline_variants":["Diagonal batching makes RMTs parallel, 3.3x faster","No retraining: Diagonal batching speeds RMTs 3.3x","Unlock RMT parallelism with diagonal batching","Diagonal batching: 3.3x speedup for long-context RMTs","RMTs go parallel: diagonal batching boosts speed 3.3x"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is the layer-level recurrent dependency graph: each (segment, layer) step may read only the same segment's previous layer and the previous segment's same-layer memory. Original Recurrent Memory Transformers violate this because the previous segment's final-layer memory enters every layer of the next segment, in which case the diagonal groups are no longer independent and the schedule is invalid.","fun_headline_variants_meta":{"raw":{"variants":["Diagonal batching makes RMTs parallel, 3.3x faster","No retraining: Diagonal batching speeds RMTs 3.3x","Unlock RMT parallelism with diagonal batching","Diagonal batching: 3.3x speedup for long-context RMTs","RMTs go parallel: diagonal batching boosts speed 3.3x"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000214,"raw_usage":{"total_tokens":1451,"prompt_tokens":994,"completion_tokens":457,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":610,"completion_tokens_details":{"reasoning_tokens":357}},"tokens_in":610,"tokens_out":457,"duration_ms":4766,"temperature":1.0,"reasoning_tokens":357,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-07T10:22:29.988168+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run one 131,072-token prompt through the same trained ARMT checkpoint twice, once with the sequential reference implementation and once with Diagonal Batching, on the same GPU, and compare final logits and wall-clock latency. If the relative logit error grows well beyond the reported roughly 2% as the number of segments increases, or if the grouped schedule is not faster than the sequential one for the claimed configurations, the exact-recurrence and speedup claims would be refuted. A direct boundary check is to apply the schedule to an original RMT with cross-layer memory flow and observe the output change, since the method is only valid for layer-level recurrent models.","supporting_citations":[{"cited_title":"Associative recurrent memory transformer.CoRR, 2024","cited_arxiv_id":null,"evidence_quote":"Defines the Associative Recurrent Memory Transformer and the broader parallel-recurrent family whose layer-segment dependency graph the method schedules."},{"cited_title":"RWKV: Reinventing RNNs for the transformer era","cited_arxiv_id":null,"evidence_quote":"Cited as another layer-level recurrent architecture whose per-layer memory would fit the diagonal schedule."}],"review_version":1}