{"id":"132c933a-3d38-4d11-b177-fc443e56fadf","arxiv_id":"2501.14577","paper_version":3,"verdict":"REJECT","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"high","formal_verification":"none","parameter_count":4,"one_line_summary":"ZETA approximates top-k attention by sorting low-dimensional keys and queries along Z-order curves, enabling parallel causal selection with O(N log N) cost.","lead":"ZETA is a new attention mechanism that finds relevant past tokens by sorting them along a one-dimensional curve, cutting the cost of Transformer attention for long sequences. It claims competitive accuracy at lower compute, but its causal-masking description is internally inconsistent and the theory is not valid for the recommended setting.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The causal-masking procedure central to ZETA is internally inconsistent: a single global sort cannot simultaneously be the order used for first-m-chunk retrieval and the original-position order needed to exclude future keys; this voids the O(N log N) parallel causal top-k claim.","rationale":"I read the paper as claiming a causal, parallel, O(N log N) top-k attention mechanism. For that claim to hold, two things must be true: (i) there is a single sorted order that can be chunked so that every query's allowable past keys are found in a bounded number of chunks, and (ii) neighbors in Z-order are near-neighbors in Euclidean space. The paper gives no proof of either; more seriously, (i) is false as stated because sorting and causality are defined over incompatible orders. The same global permutation cannot be both 'sorted by Z-value' and 'ordered by original position.' The counterexample in concrete_test makes this failure observable for a small, valid input. This is not a matter of empirical disagreement or missing baseline; it is an internal inconsistency in the core algorithm. I credit the empirical effort in Tables 1-4 and the Triton implementation, but those results do not repair the mechanism by which causal top-k search is supposed to be parallelized. The theoretical dK=3 motivation is also unsound for the stated regime, since the displayed bound contains sqrt(1−sqrt(C ln m/dK)), which is not real for dK < C ln m. If the authors can specify a correct causal chunking procedure and release code, the paper could be reconsidered, but as written the central claim is unsupported.","tokens_in":23604,"tokens_out":4813,"duration_ms":66654,"concrete_test":"Run this counterexample through Algorithm 1 as written. Let N=8, chunk size M=2, dK=1, and let the single key coordinate at original positions 0..7 be 7,6,5,4,3,2,1,0, so that after ascending Z-order sort chunk 0 contains original keys 7 and 6. For query i=2 (m=1), the only candidate chunk is chunk 0, whose keys are original positions 7 and 6, both greater than m×M−1=1; after causal exclusion the query has zero valid keys, not k=2, and attends only to the mean-history smoothing vector. Repeating for i=3..5 gives the same result. This demonstrates that the global-sort-plus-first-m-chunks scheme does not enforce causality unless an additional per-query prefix sort is introduced, and no such mechanism is described.","verdict_should_be":"UNCHANGED","load_bearing_attack":"Section 3.2.2 and Algorithm 1 state that after globally sorting Z-order keys into chunks, a query in chunk m searches only the first m chunks, 'indexing the original unsorted keys from 0 to m×M−1 in the sorted list,' to enforce causality. This is not coherent: the sort is a permutation of all N keys, so the first m sorted chunks generally contain original positions with index greater than m×M−1, and for early queries they may consist entirely of future positions. The explicit exclusion step 'j > m×M' then leaves far fewer than k valid keys, often zero, so the claimed top-k retrieval and the O(N log N) parallel training are not achieved. Making the procedure consistent by sorting each prefix separately would cost O(N^2 log N), or would require an incremental data structure that the paper does not specify, contradicting the central efficiency claim. The windowed top-k retrieval also has no proof that Z-order contiguity captures Euclidean top-k, and the theoretical bound in Theorem 3.3 is not usable at dK=3 because sqrt(1 − sqrt(C ln m / dK)) is imaginary when dK < C ln m. The causal-masking conflict alone is sufficient to invalidate the central claim.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes ZETA, a sparse top-k attention mechanism that maps keys and queries to a low-dimensional space, then to one dimension via Z-order curves, sorts the keys, and partitions them into chunks to enable parallel top-k retrieval under causal masks. It claims O(N log N) time and space complexity and reports experiments on the Multi-Query Associative Recall task, Long Range Arena, and WikiText-103, along with a theoretical analysis motivating small key/query dimensionality and an adaptive Cauchy softmax. The paper is written as an ICLR 2025 conference paper and includes a custom Triton implementation with efficiency measurements.","tokens_in":23804,"tokens_out":5950,"duration_ms":50216,"significance":"If the algorithm worked as described, ZETA would address a genuine bottleneck: causal top-k attention is difficult to parallelize because future tokens must be excluded before retrieval. The paper also contains a substantial empirical component and a custom Triton implementation with reported speedups. However, the central algorithmic mechanism is internally inconsistent (see Major Comment 1), and the theoretical bound used to justify the low-dimensional projection is invalid in the relevant regime (Major Comment 2). These issues undermine the paper's main claims, so the current manuscript cannot be accepted.","major_comments":[{"comment":"The described causal chunked search is internally inconsistent. After globally sorting the Z-order keys, the sorted list is a permutation of all N original positions. The paper says that a query in chunk m searches only the first m chunks and 'index[es] the original unsorted keys from 0 to m×M−1 in the sorted list,' and Algorithm 1 Step 4 says to 'Exclude keys from positions j > m×M.' These two conditions cannot hold simultaneously: the first m chunks of the sorted order will generally contain keys with original indices larger than m×M−1, and after excluding them, the number of valid keys is typically much less than k and often zero for early queries. Consequently, the method as specified does not retrieve k past keys per query and does not realize the claimed O(N log N) parallel causal top-k attention. Sorting each prefix separately would cost O(N^2 log N), and no incremental data structure is specified; thus the central efficiency claim is unsupported.","section":"Section 3.2.2 and Algorithm 1"},{"comment":"The risk bound in Theorem 3.3 contains the factor sqrt(1 − sqrt(C ln m / dK)). In the proof, epsilon is set to sqrt(C ln m / dK) in the Johnson–Lindenstrauss Lemma, which is only valid for epsilon < 1. For the paper's recommended dK = 3, this requires m < exp(3/C); for the sequence lengths considered in the paper (e.g., N = 4096 in LRA), m will typically exceed this threshold, making the bound imaginary and the derivation invalid. The theorem therefore does not establish the claimed trade-off between the curse of dimensionality and preservation of relative distances in the regime where ZETA operates.","section":"Section 3.2.1 and Appendix A"},{"comment":"The nearest-neighbor retrieval step selects a window of size k around the query's insertion position in the sorted Z-order list and treats the keys in that window as the top-k set. The paper supplies no bound or proof that the Euclidean top-k neighbors are contained in such a window after sorting; Z-order curves preserve locality only approximately, and the locality experiment in Section 4.4 measures neighbor overlap after projection, not contiguity after sorting. Without a retrieval-recall guarantee, the method is not demonstrably a top-k attention method, and the experimental accuracy cannot be attributed to the described selection mechanism.","section":"Section 3.2.2"}],"minor_comments":[{"comment":"The paragraph before Lemma A.1 contains an incomplete sentence and an unresolved URL: 'The recent paper, Reformer, proposed the Shared-QK Transformer... https://arxiv.org/abs/2001.0445a very simple but efficient technique.' This appears to be a leftover note and should be rewritten or removed.","section":"Appendix A"},{"comment":"There is a typo in 'he @triton.autotune decorator is used...' which should read 'The @triton.autotune decorator is used...'.","section":"Appendix D"},{"comment":"The claim of consistent improvement over baselines is qualified by the Pathfinder row: ZETA's accuracy (68.20) is lower than that of the vanilla Transformer (71.40) and several other variants. The narrative in Section 4.1 should be adjusted to reflect this.","section":"Table 2"},{"comment":"The notation is confusing: the top-k value is sometimes written as k and sometimes as K (e.g., in Section 3.2.2, 'a window of size K centered around the insertion point'). Please standardize the notation.","section":"Section 3.2.2"},{"comment":"The Limitations section acknowledges the general risk of ignoring important tokens in top-k attention, but it does not mention the causal-masking inconsistency described in Major Comment 1, which is a more direct limitation of the proposed algorithm.","section":"Appendix F"}],"recommendation":"reject","confidential_remarks":"The internal inconsistency in the central causal-search procedure is not a local presentation issue; fixing it would require a different algorithmic design or a substantial change to the claimed complexity. The theoretical bound in Theorem 3.3 is also invalid for the recommended dK=3 setting. These are load-bearing errors that cannot be repaired within the scope of the current manuscript."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"The paper attacks a real problem: parallelizing top-k attention under causal masks. The Z-order curve route is genuinely new, and the combination of low-dimensional keys/queries, chunked sorting, and a trainable Cauchy softmax is not in Reformer, IceFormer, or earlier top-k work. If the numbers are faithful, ZETA is competitive on LRA and roughly on par with vanilla on WikiText, and the Triton implementation shows real speedups. That is worth acknowledging.\n\nThe soft spots are not minor. The causal-masking logic in Section 3.2.2 is internally inconsistent. A single global sort cannot be both the sorted order used for \"first m chunks\" retrieval and the original-position order used to exclude future keys. For early queries, the first sorted chunks may contain mostly future positions, so the causal filter leaves few or no valid keys. Sorting each prefix separately would cost O(N^2 log N), contradicting the headline O(N log N). The paper never specifies a structure that reconciles the two orders.\n\nTheorem 3.3's bound contains sqrt(1 - sqrt(C ln m / dK)). With dK=3, the inner term is larger than 1 for any realistic m, making the bound imaginary. That is the exact setting the paper adopts, so the theory does not support the low-dimensional choice; it actively breaks there. The abstract also overstates the WikiText result: Table 1 shows ZETA at 26.3 versus vanilla's 26.2, essentially a tie, while CosFormer is clearly better. And there is no released code, which matters for a submission whose value is empirical.\n\nThe central efficiency and correctness claims are therefore not established as written. I suspect the authors had a working implementation and a garbled write-up; the empirical numbers are plausible. But a reader cannot verify the algorithm from the text, and the theory is unusable in the operating regime. This is a case where a promising idea needs a serious rewrite rather than a desk rejection. I would send it to referees, because the novelty and experimental scope merit close scrutiny, but I would not build on it until the algorithm is described coherently, code is released, and the bound is repaired or removed.","headline":"The Z-order idea is fresh and the experiments are real, but the causal-masking procedure as written is self-contradictory and the main theoretical bound goes imaginary in exactly the regime the paper uses, so the central claims don't hold up.","tokens_in":24418,"tokens_out":4574,"would_cite":false,"duration_ms":41695,"reading_group":"maybe","serious_thinker":"no","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"This paper claims that top-k attention can be made parallel across entire sequences under causal masks by sorting Z-order projections of low-dimensional keys and queries, and that the resulting method matches or outperforms standard…","keywords":["top-k attention","Z-order curve","Morton code","causal masking","sparse attention","long-sequence transformers","Cauchy kernel","near-neighbor search"],"falsifier":"Run ZETA on a trained checkpoint, compute each query's exact $k$ nearest keys by Euclidean distance on the same low-dimensional keys, and measure the recall of ZETA's selected indices against that exact set; if recall is low on real data, the windowed-contiguity assumption fails.","tokens_in":23299,"feed_emoji":"⚡","tokens_out":9853,"duration_ms":79178,"temperature":0.7,"pith_summary":"ZETA sets out to remove a bottleneck in top-$k$ attention: with a causal mask, every query may only attend to past tokens, and earlier methods either mask after selection (leaving early queries with nothing to attend to) or walk through the sequence token by token, losing parallelism. The paper's solution is to project keys and queries into a low-dimensional space, map them to one-dimensional Z-order codes, sort the codes once, and let each query binary-search for its own code and read a fixed window of neighbours from the sorted prefix. If this works, training and inference for sparse attention run in $O(N\\log N)$ time and memory while keeping the whole sequence parallel on accelerators. The paper backs this with a trade-off analysis for key/query dimensionality, a trainable Cauchy softmax over Euclidean distances, and empirical results where ZETA matches vanilla attention on associative recall, posts the best average Long Range Arena accuracy among the compared models, and reaches near-Transformer perplexity on WikiText-103.","feed_headline":"Z-order curves let top-k attention go parallel","feed_subtitle":"One sort turns causal sparse attention into a windowed neighbor search that scales to long sequences.","key_machinery":"The load-bearing object is the Z-order curve (Morton code), which maps a $d$-dimensional point to a one-dimensional integer by interleaving the bits of its coordinates so that nearby points tend to get nearby codes. ZETA sorts the key codes once, splits the sorted list into chunks, and for each query in chunk $m$ binary-searches the query's code inside the first $m$ chunks and reads a fixed-size window around that insertion position as the top-$k$ set. The Adaptive Cauchy-Softmax then converts Euclidean distances in the low-dimensional space into attention weights through $\\frac{1}{\\|q-k\\|^2 + \\gamma^2}$, with $\\gamma$ trainable per layer.","core_discovery":"On its own terms, ZETA's central claim is that the top-$k$ attended tokens for every query in a causally masked sequence can be found in parallel by sorting all keys once in Z-order space, splitting the sorted list into chunks, and restricting each query in chunk $m$ to keys in the first $m$ chunks; the query's insertion position in that prefix locates a window of $k$ nearest neighbours. To make this geometry meaningful, the paper argues that keys and queries should have much lower dimension than values: small $d_K$ preserves relative distances after Z-order projection (supported by a bound built on the Johnson–Lindenstrauss lemma and a Lipschitz regression risk analysis), while a large $d_V$ preserves semantic richness. Attention weights are then computed with an Adaptive Cauchy-Softmax over Euclidean distances, with a trainable $\\gamma$ controlling the receptive field, and the paper reports that ZETA matches vanilla attention on Multi-Query Associative Recall, achieves the best average accuracy among the compared models on Long Range Arena, reaches 26.3 perplexity on WikiText-103 against 26.2 for vanilla attention, and runs faster than a FlashAttention baseline at long sequence lengths.","pith_inferences":["Beyond the paper, a natural stress test is to measure recall@k of ZETA's window retrieval against exact kNN on real trained keys; if recall is low, a data-dependent window size or a candidate re-ranking stage would be a direct fix.","The chunk-prefix causal rule may discard past keys whose Z-order codes land in later chunks; an alternative is to sort each prefix independently or use a hierarchical index, trading speed for exactness.","If the low-dimensional matching insight generalizes, other distance-based sparse attention schemes (LSH, product quantization, locality-sensitive retrieval) could adopt the same split between small matching dimensions and large value dimensions.","Because the Cauchy kernel's $\\gamma$ is trained per layer, one could test whether learned receptive fields correlate with task structure: sharp in local tasks, broad in long-range retrieval tasks."],"forward_implications":["A causal top-$k$ attention layer can be trained in parallel over the full sequence, with sorting dominating the cost at $O(N\\log N)$ time and space rather than the $O(N^2)$ of standard attention.","Key and query dimensions can be set far below value dimensions ($d_K = d_Q = 3$ in the paper's experiments) without degrading accuracy, because matching needs relative-distance preservation rather than semantic richness.","Replacing the exponential softmax with a trainable Cauchy kernel over Euclidean distances keeps long-range tokens influential while letting each layer adjust its receptive field.","On the reported benchmarks ZETA matches vanilla attention on Multi-Query Associative Recall, posts the best average LRA accuracy among the compared models, and reaches 26.3 WikiText-103 perplexity versus 26.2 for vanilla attention.","At sequence lengths from 8K to 64K, the paper's custom-kernel implementation is faster than FlashAttention in forward and forward-backward passes, with modestly higher memory use."],"supporting_citations":[{"why":"Defines the standard attention and dot-product softmax that ZETA replaces.","marker":"(Vaswani et al., 2017)"},{"why":"Introduces the top-k attention formulation that ZETA extends to causal parallel training.","marker":"(Gupta et al., 2021)"},{"why":"Supplies the shared key-query projection assumption and an earlier top-k style efficient attention.","marker":"(Kitaev et al., 2020)"},{"why":"Documents the per-token sequential bottleneck in causal kNN attention that ZETA aims to remove.","marker":"(Mao et al., 2024)"},{"why":"Gives the Johnson-Lindenstrauss lemma used to justify low-dimensional keys and queries.","marker":"(Johnson et al., 1986)"},{"why":"Defines the Long Range Arena benchmark and the hyperparameter settings used in evaluation.","marker":"(Tay et al., 2021)"},{"why":"Defines the Multi-Query Associative Recall task used as the synthetic recall experiment.","marker":"(Arora et al., 2024a)"},{"why":"Provides the WikiText-103 dataset and perplexity evaluation.","marker":"(Merity et al., 2017)"},{"why":"Supplies the FlashAttention baseline that ZETA's efficiency experiments compare against.","marker":"(Dao et al., 2022; Dao, 2024)"},{"why":"Is the cited source for Z-order curves and their locality-preserving one-dimensional mapping.","marker":"(Dugundji, 1966)"}],"fun_headline_variants":["ZETA: parallel top-k attention via Z-order sorting","Z-order curves make causal top-k attention parallel","One sort enables parallel top-k attention for long sequences","ZETA: low-dim keys, Z-order sort, parallel top-k attention"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The central assumption is that the true top-$k$ keys for a query are contiguous around the query's insertion position in the sorted Z-order list, so a fixed window over the first $m$ chunks recovers them without scanning the full prefix.","fun_headline_variants_meta":{"raw":{"variants":["ZETA: parallel top-k attention via Z-order sorting","Z-order curves make causal top-k attention parallel","One sort enables parallel top-k attention for long sequences","ZETA: low-dim keys, Z-order sort, parallel top-k attention"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000192,"raw_usage":{"total_tokens":1426,"prompt_tokens":1102,"completion_tokens":324,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":718,"completion_tokens_details":{"reasoning_tokens":256}},"tokens_in":718,"tokens_out":324,"duration_ms":3569,"temperature":1.0,"reasoning_tokens":256,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-10T15:01:12.000129+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run ZETA on a trained checkpoint, compute each query's exact $k$ nearest keys by Euclidean distance on the same low-dimensional keys, and measure the recall of ZETA's selected indices against that exact set; if recall is low on real data, the windowed-contiguity assumption fails.","supporting_citations":[{"cited_title":"Gomez, ukasz Kaiser, and Illia Polosukhin","cited_arxiv_id":null,"evidence_quote":"Defines the standard attention and dot-product softmax that ZETA replaces."},{"cited_title":"Extensions of lipschitz maps into banach spaces","cited_arxiv_id":null,"evidence_quote":"Gives the Johnson-Lindenstrauss lemma used to justify low-dimensional keys and queries."},{"cited_title":"Topology","cited_arxiv_id":null,"evidence_quote":"Is the cited source for Z-order curves and their locality-preserving one-dimensional mapping."}],"review_version":1}