Pith. sign in

REVIEW 4 major objections 4 minor 67 references

Adjoint sharding for very long context training of state space models

T0 review · 4 major / 4 minor · reviewed 2026-08-10 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read Adjoint sharding computes backpropagation-equivalent gradients for state-space models by decomposing them into independent vector-Jacobian products, cutting training memory by up to 3X at 1M-token context.

desk verdict A clean adjoint-method decomposition for SSM gradients with a real parallelization idea, but the abstract's memory-reduction claims are unsupported by any experiments and the paper's own storage analysis undermines them. read the letter →

arxiv 2501.00692 v1 pith:55GVQLEE submitted 2025-01-01 cs.LG cs.AIcs.CL

Xingzi Xu , Amir Tavanaei , Kavosh Asadi , Karim Bouyarmane This is my paper · ORCID
classification cs.LGcs.AIcs.CL
keywords adjointmethodstate-spacemodelslong-contexttraininggradientshardingvector-Jacobianproductsmemory-efficienttruncatedbackpropagationthroughtimedistributed
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

Training large language models on very long contexts is bottlenecked by GPU memory, because backpropagation stores activations for the whole sequence, so memory grows rapidly with context length. This paper proposes adjoint sharding, which replaces backpropagation with the adjoint method for the recurrent state-space layers, decomposing the gradient into many independent vector-Jacobian products (one per token per layer) that can be computed, stored, and discarded separately. The paper proves that summing these products reproduces exactly the backpropagation gradient, and argues that sharding them across time and layers reduces memory enough to train a 1.27B-parameter state-space model on contexts above 100K tokens with a small cluster. A truncated variant cuts the number of vector-Jacobian products from quadratic to linear in context length. If the memory reduction holds, long-context training and fine-tuning of state-space architectures becomes tractable on far fewer GPUs than today.

What carries the argument

The load-bearing object is the adjoint state $\lambda_{t,\tau}^k = C_t^k \prod_{i=1}^{t-\tau} A_{t+1-i}^k$ for layer $k$, which replaces the chained derivatives of backpropagation. Proposition 1 (the adjoint method for recurrences) supplies the recurrence $\lambda_{i-1}=\lambda_i\,\partial f(i,h_{i-1},\theta)/\partial h_{i-1}$ that makes these states computable by a backward pass without storing the forward computation graph. The residual identity $y_t^K = y_t^0 + \sum_k \tilde{y}_t^k$ then lets the gradient separate into a sum over layers. Each term in the resulting sum is a vector-Jacobian product of the form $\mathrm{vjp}(v) = v \cdot \partial \mathrm{Net}_\theta(\mathrm{Input}_i)$, which can be evaluated with reverse-mode automatic differentiation initialized with the vector $v$, and which the paper treats as costing about one forward pass. The independence of these VJPs is what allows memory to be sharded and computation to be parallelized.

What would settle it

Train a 1.27B-parameter state-space model on a 1M-token sequence with adjoint sharding and record peak GPU memory, comparing with backpropagation under activation checkpointing on the same hardware; if adjoint sharding's peak memory is not roughly 3X lower, the abstract's central memory claim is contradicted.

Watch

Extended reading notes

Core claim

The central claim is that for a residual stack of state-space layers, the gradient of the training loss with respect to every parameter can be written as a sum (over tokens $t$ and layers $k$) of independent vector-Jacobian products, with no cross-dependencies between terms. For a single SSM layer, Proposition 2 gives $$\frac{d l_t}{d\$\theta$} = \Big[\sum_{i=1}^t \mathrm{vjp}^{A_i}\big(\tfrac{dl_t}{dy_t}\lambda_{t,i}\otimes h_{i-1}\big)\Big] \oplus \Big[\sum_{i=1}^t \mathrm{vjp}^{B_i}\big(\tfrac{dl_t}{dy_t}\lambda_{t,i}\otimes \hat{x}_i\big)\Big] \oplus \mathrm{vjp}^{C_t}\big(\tfrac{dl_t}{dy_t}\otimes h_t\big),$$ with adjoint state $\lambda_{t,\tau}=C_t\prod_{i=1}^{t-\tau} A_{t+1-i}$; Proposition 3 extends this to $K$ stacked residual layers, so the full gradient is a sum of such terms over $t$ and $k$. Because each vector-Jacobian product depends only on the saved forward states and the adjoint state, the products can be sharded by time and layer across devices, computed in parallel, and discarded immediately. The paper also introduces truncated adjoint sharding, which limits the time sum to the last $\bar{T}$ states, cutting the number of VJPs from $O(T^2)$ to $O(T\bar{T})$ while retaining implicit dependence on all earlier states through the recurrence.

Load-bearing premise

The promised memory savings depend on the assumption that storing the intermediate values the method needs for every token and every layer—the transition and output matrices, hidden states, and normalized inputs—costs less memory than backpropagation's stored activations, so peak memory still drops sharply.

Editorial extensions

If this is right

  • Training or fine-tuning a 1.27B-parameter SSM on contexts above 100K tokens becomes possible on a small cluster of five instances, where backpropagation tops out near 35K tokens.
  • Gradient computation for residual and recurrent models becomes embarrassingly parallel across time and layers, so wall-clock training time can drop if the independent VJPs are scheduled across GPUs or GPU instances.
  • Truncated adjoint sharding computes an approximate gradient in linear time, making very long contexts computationally feasible where the full method's $O(T^2)$ VJP count would be prohibitive.
  • Because the gradient decomposition is exact for the full method, training with adjoint sharding reproduces backpropagation's training results while using less memory.
  • Activations, gradients, and optimizer states can all be sharded across devices, so per-GPU memory falls roughly as $1/\Upsilon$ with $\Upsilon$ devices.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • The paper does not report measured memory for the 1M-token, 1.27B-parameter case its abstract cites, only smaller-scale curves; whether the stored states keep peak memory below backpropagation at that scale is an open empirical question.
  • If the VJP decomposition is computed in FP16, rounding may accumulate differently than in backpropagation's sequential chain, so 'equivalent gradients' likely holds only up to numerical precision; a mixed-precision gradient-equality check would clarify this.
  • Truncated adjoint sharding is effectively a truncated-backpropagation-through-time variant for SSMs; its convergence and the optimal truncation length $\bar{T}$ are left for future work, and likely depend on how fast the products $\prod A_t$ decay, which is a testable prediction.
  • The paper's parallel-speedup estimate assumes VJPs are compute-bound at roughly one forward pass each; if memory bandwidth is the binding constraint, the practical speedup on GPU instances will fall below the analyzed figure.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

4 major / 4 minor

Summary. The paper proposes adjoint sharding, a method that computes gradients for state-space model (SSM) language models by decomposing backpropagation into independent vector-Jacobian products via the adjoint method. It claims this reduces training memory by up to 3X (abstract) or "orders of magnitude" (abstract and conclusion), enabling training on 1M-token contexts for a 1.27B parameter model on five AWS P4 instances. The paper derives gradient-equivalence results in Propositions 2 and 3, introduces a truncated variant and a distributed algorithm, and presents an analytical memory/FLOP analysis. The manuscript contains no experimental section, no measured memory numbers, and no code.

Significance. If the memory-reduction claims were established, adjoint sharding would be a practically valuable technique for long-context SSM training. The gradient-equivalence derivations in Propositions 2 and 3 appear internally consistent and provide a useful formalization of the adjoint approach for recurrent SSM layers. However, the practical impact is entirely unverified: the paper's central empirical assertions are unsupported, and its own memory accounting raises serious doubts about the claimed savings. For a cs.LG systems submission, the absence of experiments is a fundamental gap.

major comments (4)
  1. [Abstract and Conclusion] The abstract states "Empirical results show the proposed adjoint sharding algorithm reduces memory usage by up to 3X with a 1.27B parameter large language model on 1M context length training," and the conclusion repeats "orders of magnitude of memory reduction," yet the manuscript contains no experimental section, no measurement protocol, no memory numbers, and no code. Figure 1 is a plot without axes or experimental setup, and Figure 6 is based on assumed speedups. This is load-bearing because the paper's contribution is a training method whose benefit is memory reduction; the gradient-equivalence proofs alone do not establish any memory savings.
  2. [Section 4.5] The memory analysis states that storing {A_t^k, C_t^k, h_t^k, y_t^k} over T tokens and K layers costs T K(2N+P) + T P FP16 numbers. For the claimed scale (T=1M, K=100, N=225, P=128), this is approximately 5.8e10 FP16 values, or about 116 GB, which exceeds the 40 GB memory of a single AWS P4 (A100) instance. Even distributed over five GPUs, this is roughly 23 GB per GPU before vjp workspaces, model weights, and optimizer states are counted. The paper never compares this peak memory with a backpropagation baseline using activation checkpointing, so the "up to 3X" and "orders of magnitude" claims are unsupported.
  3. [Algorithm 1] Algorithm 1, line 10, requires storing per-token, per-layer states {h_t^k, C_t^k, y_hat_t^k, A_t^k} on each device. This storage grows as O(T K N) with context length and layer count and is not a constant-memory scheme. The claimed advantage of adjoint sharding over backpropagation therefore depends on comparing this O(T K N) state storage against the memory cost of the autograd graph in BPTT, a comparison the paper does not provide. The theoretical gradient equivalence does not imply lower memory.
  4. [Section 4.3] Truncated adjoint sharding is introduced with the phrase "empirically showcase its performance," but no empirical evaluation appears anywhere in the manuscript. The same section later states "We leave investigation of bar(T)'s impact on performances for future works," which explicitly acknowledges that the performance-maintenance claim is not established. The statement that truncated adjoint sharding achieves "similar results" is therefore unsupported.
minor comments (4)
  1. [Abstract and Conclusion] The abstract reports "up to 3X" memory reduction, while both the abstract's earlier sentence and the conclusion claim "orders of magnitude" reduction; these are inconsistent quantitative claims.
  2. [Appendix A.2] In the proof of Proposition 3, the vjpBi terms are written with x_hat_i^k and later changed to y_hat_i^{k-1}; the notation should be harmonized to match Proposition 3 and the main text.
  3. [Figures 1 and 6] Figure 1 and Figure 6 lack axis labels, units, and any description of the experimental or simulated setup, so they cannot be interpreted or reproduced.
  4. [References] The bibliography contains duplicate entries: reference [4] is repeated, and references [52] and [53] both cite the ZeRO paper; these should be consolidated.

Circularity Check

0 steps flagged · score 2.0 of 10

No load-bearing circularity: the gradient derivation is a direct chain-rule expansion, and the only self-citation is non-load-bearing.

full rationale

The derivative chain is self-contained. Proposition 1 is the standard adjoint-sensitivity formula quoted from the independent references [8,32]; Proposition 2 applies it to the linear SSM recurrence ht = At ht-1 + Bt xhat_t, yt = Ct ht by direct substitution (Eq. 8), and Proposition 3 sums over the residual layers. These are rearrangements of the same chain-rule terms that backpropagation computes, so the claimed gradient equivalence is derived, not assumed. No parameter is fitted to a subset of data and then reported as a prediction: the truncation window \(\bar{T}\) is a user-set hyperparameter whose convergence the paper explicitly leaves to future work, and no experimental memory numbers are used in the derivation. The only author-overlapping citation is [66] in the background sentence "The adjoint method is a constant-memory optimization technique for dynamical systems [9, 66]"; it is paired with the independent [9] and is not load-bearing. The abstract's "up to 3X" and "orders of magnitude" memory claims are unsupported by an experimental section, but that is a verifiability/correctness gap, not circularity, and the paper itself flags the missing implementation and convergence analysis in its Limitation paragraph.

Assumptions & free parameters 1 free parameters · 4 assumptions · 0 invented entities

The central derivation uses standard adjoint equations and the assumed SSM residual structure. No invented entities appear. The only free parameter is the truncation window \bar{T}, whose value is not fitted in the paper.

free parameters (1)
  • Truncation window \bar{T} = not specified
    The truncation window \bar{T} controls the trade-off between speed and approximation quality in truncated adjoint sharding. Its value is chosen by the user and is not derived from theory or fixed by data in the paper.
assumptions (4)
  • domain assumption SSM dynamics h_t = A_t h_{t-1} + B_t x_t and y_t = C_t h_t.
    Section 3.1 defines the model class that the adjoint derivation uses; the gradient formulas depend on this linear recurrent form.
  • domain assumption Residual structure y_t^K = y_t^0 + sum_k \tilde{y}_t^k.
    Section 3.2 introduces the residual network; Proposition 3 uses it to separate gradients across layers.
  • standard math Adjoint equations from Proposition 1 of [8] are correct.
    Section 3.3 states the adjoint method for recurrence relations as a cited proposition and builds on it without proof.
  • domain assumption VJPs can be computed with memory comparable to a forward pass and can be parallelized on modern GPUs.
    Section 4.5 and the Limitations paragraph rely on this implementation assumption to argue for speedups, but provide no measured implementation.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Adjoint sharding for very long context training of state space models." pith.science (2026). https://pith.science/paper/55GVQLEE

@misc{pith2026250100692,
  author       = {Pith},
  title        = {Pith review of: Adjoint sharding for very long context training of state space models},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/55GVQLEE}},
  note         = {Machine review of arXiv:2501.00692}
}
read the original abstract

Despite very fast progress, efficiently training large language models (LLMs) in very long contexts remains challenging. Existing methods fall back to training LLMs with short contexts (a maximum of a few thousands tokens in training) and use inference time techniques when evaluating on long contexts (above 1M tokens context window at inference). As opposed to long-context-inference, training on very long context input prompts is quickly limited by GPU memory availability and by the prohibitively long training times it requires on state-of-the-art hardware. Meanwhile, many real-life applications require not only inference but also training/fine-tuning with long context on specific tasks. Such applications include, for example, augmenting the context with various sources of raw reference information for fact extraction, fact summarization, or fact reconciliation tasks. We propose adjoint sharding, a novel technique that comprises sharding gradient calculation during training to reduce memory requirements by orders of magnitude, making training on very long context computationally tractable. Adjoint sharding is based on the adjoint method and computes equivalent gradients to backpropagation. We also propose truncated adjoint sharding to speed up the algorithm while maintaining performance. We provide a distributed version, and a paralleled version of adjoint sharding to further speed up training. Empirical results show the proposed adjoint sharding algorithm reduces memory usage by up to 3X with a 1.27B parameter large language model on 1M context length training. This allows to increase the maximum context length during training or fine-tuning of a 1.27B parameter model from 35K tokens to above 100K tokens on a training infrastructure composed of five AWS P4 instances.

Figures

Figures reproduced from arXiv: 2501.00692 by the authors.

Figure 1
Figure 1. Compared to backpropagation (red lines), adjoint sharding (blue lines) significantly reduces memory re [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Adjoint sharding dissembles large models’ gra [PITH_FULL_IMAGE:figures/full_fig_p002_2.png] view at source ↗
Figure 3
Figure 3. Lines in red are fine-tuning free methods and lines in blue are fine-tuning methods. Fine-tuning methods [PITH_FULL_IMAGE:figures/full_fig_p004_3.png] view at source ↗
Figures from the paper (3 more)
Figure 4
Figure 4. Figure 4: The adjoint states are computed sequentially backwards. [PITH_FULL_IMAGE:figures/full_fig_p006_4.png]
Figure 5
Figure 5. Figure 5: Computation schematic of dl t/dθAk , dl t/dθBk , and dl t/dθCk . We provide the proof to proposition 3 in section A.2. Define Λ t k = {λ t,τ k } t τ=1, proposition 3 shows that the gradients of each network’s parameters computed with each token only correlate through t…
Figure 6
Figure 6. Figure 6: Training time (/day) per epoch comparison for [PITH_FULL_IMAGE:figures/full_fig_p010_6.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

67 extracted references · 19 canonical work pages

  1. [1]

    Blackmamba: Mixture of experts for state-space models, 2024

    Quentin Anthony, Yury Tokpanov, Paolo Glorioso, and Beren Millidge. Blackmamba: Mixture of experts for state-space models, 2024. URL https://arxiv.org/abs/2402.01771

    arXiv 2024

  2. [2]

    Fast jacobian-vector product for deep networks, 2021

    Randall Balestriero and Richard Baraniuk. Fast jacobian-vector product for deep networks, 2021. URL https: //arxiv.org/abs/2104.00219

    arXiv 2021

  3. [4]

    Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind

    Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey, 2018. URL https://arxiv.org/abs/1502.05767

    arXiv 2018

  4. [5]

    xlstm: Extended long short-term memory, 2024

    Maximilian Beck, Korbinian P ¨oppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, G¨unter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory, 2024. URL https://arxiv.org/abs/2405.04517

    arXiv 2024

  5. [6]

    Peters, and Arman Cohan

    Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer, 2020. URL https://arxiv.org/abs/2004.05150

    arXiv 2020

  6. [7]

    Internlm2 technical report,

    Zheng Cai, Maosong Cao, Haojiong Chen, Kai Chen, Keyu Chen, Xin Chen, et al. Internlm2 technical report,

    work page

  7. [8]

    Adjoint sensitivity analysis for differential-algebraic equations: algorithms and software

    Yang Cao, Shengtai Li, and Linda Petzold. Adjoint sensitivity analysis for differential-algebraic equations: algorithms and software. Journal of Computational and Applied Mathematics , 149(1):171–191, 2002. ISSN 0377-0427. doi: https://doi.org/10.1016/S0377-0427(02)00528-9. URL https://www.sciencedirect.com/ science/article/pii/S0377042702005289. Scientif...

    work page doi:10.1016/s0377-0427(02)00528-9 2002

  8. [9]

    Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equa- tions, 2019. URL https://arxiv.org/abs/1806.07366

    arXiv 2019

Show all 67 references
  1. [10]

    Extending context window of large language models via positional interpolation, 2023

    Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation, 2023. URL https://arxiv.org/abs/2306.15595

    2023 arXiv

  2. [11]

    Longlora: Efficient fine-tuning of long-context large language models, 2024

    Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models, 2024. URL https://arxiv.org/abs/2309.12307

    2024 arXiv

  3. [12]

    The backpropagation algorithm for a math student, 2023

    Saeed Damadi, Golnaz Moharrer, and Mostafa Cham. The backpropagation algorithm for a math student, 2023. URL https://arxiv.org/abs/2301.09977

    2023 arXiv

  4. [13]

    Flashattention-2: Faster attention with better parallelism and work partitioning, 2023

    Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning, 2023. URL https: //arxiv.org/abs/2307.08691

    2023 arXiv

  5. [15]

    Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024

    Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024. URL https://arxiv.org/abs/2405.21060

    2024 arXiv

  6. [16]

    Fu, Stefano Ermon, Atri Rudra, and Christopher R´e

    Tri Dao, Daniel Y . Fu, Stefano Ermon, Atri Rudra, and Christopher R´e. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022. URL https://arxiv.org/abs/2205.14135

    2022 arXiv

  7. [17]

    Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, and Caglar Gulc...

    2024

  8. [19]

    Longrope: Extending llm context window beyond 2 million tokens, 2024

    Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending llm context window beyond 2 million tokens, 2024. URL https://arxiv.org/ abs/2402.13753

    2024 arXiv

  9. [20]

    Augmented neural odes, 2019

    Emilien Dupont, Arnaud Doucet, and Yee Whye Teh. Augmented neural odes, 2019. URL https://arxiv. org/abs/1904.01681. 11 PREPRINT - D ECEMBER 31, 2024

    2019 arXiv

  10. [21]

    Fu, Tri Dao, Khaled K

    Daniel Y . Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher R ´e. Hungry hungry hippos: Towards language modeling with state space models, 2023. URL https://arxiv.org/abs/2212. 14052

    2023

  11. [23]

    Mamba: Linear-time sequence modeling with selective state spaces, 2024

    Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024. URL https: //arxiv.org/abs/2312.00752

    2024 arXiv

  12. [24]

    Combining recurrent, convolutional, and continuous-time models with linear state-space layers, 2021

    Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher R ´e. Combining recurrent, convolutional, and continuous-time models with linear state-space layers, 2021. URL https:// arxiv.org/abs/2110.13985

    2021 arXiv

  13. [25]

    How to train your hippo: State space models with generalized orthogonal basis projections, 2022

    Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher R´e. How to train your hippo: State space models with generalized orthogonal basis projections, 2022. URL https://arxiv.org/abs/2206.12037

    2022 arXiv

  14. [26]

    Attention mechanisms in computer vision: A survey

    Meng-Hao Guo, Tian-Xing Xu, Jiang-Jiang Liu, Zheng-Ning Liu, Peng-Tao Jiang, Tai-Jiang Mu, Song-Hai Zhang, Ralph R Martin, Ming-Ming Cheng, and Shi-Min Hu. Attention mechanisms in computer vision: A survey. Computational visual media, 8(3):331–368, 2022

    2022

  15. [27]

    Simplifying and understanding state space models with diago- nal linear rnns, 2023

    Ankit Gupta, Harsh Mehta, and Jonathan Berant. Simplifying and understanding state space models with diago- nal linear rnns, 2023. URL https://arxiv.org/abs/2212.00768

    2023 arXiv

  16. [28]

    Deep residual learning for image recognition, 2015

    Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015. URL https://arxiv.org/abs/1512.03385

    2015 arXiv

  17. [29]

    Deep residual learning for image recognition

    Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016

    2016

  18. [30]

    Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory, 2019

    Julien Herrmann, Olivier Beaumont, Lionel Eyraud-Dubois, Julien Hermann, Alexis Joly, and Alena Shilova. Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory, 2019. URL https://arxiv.org/abs/1911.13214

    2019

  19. [31]

    A tutorial on training recurrent neural networks , covering bppt , rtrl , ekf and the ” echo state network ” approach - semantic scholar

    Herbert Jaeger. A tutorial on training recurrent neural networks , covering bppt , rtrl , ekf and the ” echo state network ” approach - semantic scholar. In National Research Center for Information Technology, 2002 , 2005. URL https://api.semanticscholar.org/CorpusID:192593367

    2002

  20. [32]

    Adjoint methods and sensitivity analysis for recurrence, 01 2007

    Steven Johnson. Adjoint methods and sensitivity analysis for recurrence, 01 2007

    2007

  21. [33]

    Linear dynamical systems as a core computational primitive

    Shiva Kaul. Linear dynamical systems as a core computational primitive. In H. Larochelle, M. Ranzato, R. Had- sell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 16808–16820. Curran Associates, Inc., 2020. URL https://pr...

    2020

  22. [34]

    Segment anything

    Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al. Segment anything. In Proceedings of the IEEE/CVF Interna- tional Conference on Computer Vision, pages 4015–4026, 2023

    2023

  23. [35]

    Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang

    Dacheng Li*, Rulin Shao*, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph E. Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang. How long can open-source llms truly promise on context length?, June 2023. URL https://lmsys.org/blog/2023-06-29-longchat

    2023

  24. [36]

    Long-context llms struggle with long in-context learning, 2024

    Tianle Li, Ge Zhang, Quy Duc Do, Xiang Yue, and Wenhu Chen. Long-context llms struggle with long in-context learning, 2024. URL https://arxiv.org/abs/2404.02060

    2024 arXiv

  25. [37]

    Jamba: A hybrid transformer-mamba language model, 2024

    Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Ere...

    2024 arXiv

  26. [38]

    Ring attention with blockwise transformers for near-infinite context,

    Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context,

  27. [39]

    World model on million-length video and language with blockwise ringattention, 2024

    Hao Liu, Wilson Yan, Matei Zaharia, and Pieter Abbeel. World model on million-length video and language with blockwise ringattention, 2024. URL https://arxiv.org/abs/2402.08268

    2024 arXiv

  28. [40]

    The llama 3 herd of models, 2024

    Meta et al. The llama 3 herd of models, 2024. URL https://arxiv.org/abs/2407.21783

    2024 arXiv

  29. [41]

    Mixed precision training, 2018

    Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training, 2018. URL https://arxiv.org/abs/1710.03740. 12 PREPRINT - D ECEMBER 31, 2024

    2018 arXiv

  30. [42]

    Schoenholz

    Roman Novak, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Fast finite width neural tangent kernel, 2022. URL https://arxiv.org/abs/2206.08720

    2022 arXiv

  31. [43]

    Matrix multiplication background user’s guide, 2024

    NVIDIA. Matrix multiplication background user’s guide, 2024. URL https://docs.nvidia.com/ deeplearning/performance/dl-performance-matrix-multiplication/index.html

    2024

  32. [44]

    Gpt-4 technical report, 2024

    OpenAI et al. Gpt-4 technical report, 2024. URL https://arxiv.org/abs/2303.08774

    2024 arXiv

  33. [45]

    Resurrecting recurrent neural networks for long sequences, 2023

    Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences, 2023. URL https://arxiv.org/abs/2303. 06349

    2023

  34. [46]

    On the difficulty of training recurrent neural networks,

    Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks,

  35. [47]

    Pytorch: An imperative style, high-performance deep learning library, 2019

    Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas K ¨opf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu F...

    2019 arXiv

  36. [48]

    Scalable diffusion models with transformers

    William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4195–4205, 2023

    2023

  37. [49]

    Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu

    Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV , Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Kri...

    2023 arXiv

  38. [50]

    Yarn: Efficient context window extension of large language models, 2023

    Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models, 2023. URL https://arxiv.org/abs/2309.00071

    2023 arXiv

  39. [51]

    Moe-mamba: Efficient selective state space models with mixture of experts, 2024

    Maciej Pi ´oro, Kamil Ciebiera, Krystian Kr ´ol, Jan Ludziejewski, Michał Krutul, Jakub Krajewski, Szymon An- toniak, Piotr Miło´s, Marek Cygan, and Sebastian Jaszczur. Moe-mamba: Efficient selective state space models with mixture of experts, 2024. URL https://arxiv.org/abs/2...

    2024 arXiv

  40. [53]

    Zero: Memory optimizations toward training trillion parameter models, 2020

    Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models, 2020. URL https://arxiv.org/abs/1910.02054

    2020 arXiv

  41. [54]

    Zero-offload: Democratizing billion-scale model training, 2021

    Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. Zero-offload: Democratizing billion-scale model training, 2021. URL https: //arxiv.org/abs/2101.06840

    2021 arXiv

  42. [55]

    Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024

    Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024. URL https://arxiv.org/abs/2407. 08608

    2024

  43. [56]

    Sohoni, Christopher R

    Nimit S. Sohoni, Christopher R. Aberger, Megan Leszczynski, Jian Zhang, and Christopher R ´e. Low-memory neural network training: A technical report, 2022. URL https://arxiv.org/abs/1904.10631

    2022 arXiv

  44. [57]

    Unbiasing truncated backpropagation through time, 2017

    Corentin Tallec and Yann Ollivier. Unbiasing truncated backpropagation through time, 2017. URL https: //arxiv.org/abs/1705.08209

    2017 arXiv

  45. [58]

    Focused transformer: Contrastive training for context scaling, 2023

    Szymon Tworkowski, Konrad Staniszewski, Mikołaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Miło ´s. Focused transformer: Contrastive training for context scaling, 2023. URL https://arxiv.org/abs/2307. 03170

    2023

  46. [59]

    Ntk-aware scaled rope, 2023

    Reddit users. Ntk-aware scaled rope, 2023. URL https://www.reddit.com/r/LocalLLaMA/comments/ 14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/

    2023

  47. [60]

    Gomez, Lukasz Kaiser, and Illia Polosukhin

    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023. URL https://arxiv.org/abs/1706.03762

    2023 arXiv

  48. [61]

    Rellermeyer

    Joost Verbraeken, Matthijs Wolting, Jonathan Katzy, Jeroen Kloppenburg, Tim Verbelen, and Jan S. Rellermeyer. A survey on distributed machine learning. ACM Computing Surveys, 53(2):1–33, March 2020. ISSN 1557-7341. doi: 10.1145/3377454. URL http://dx.doi.org/10.1145/3377454. 1...

    2020 doi

  49. [62]

    An empirical study of mamba-based language models, 2024

    Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shoeybi, and Bryan Catanzaro. An empirical study of mamba-...

    2024 arXiv

  50. [63]

    State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory, 2023

    Shida Wang and Beichen Xue. State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory, 2023. URL https://arxiv.org/abs/2309.13414

    2023 arXiv

  51. [64]

    P.J. Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE , 78(10): 1550–1560, 1990. doi: 10.1109/5.58337

    1990 doi

  52. [65]

    Efficient streaming language models with attention sinks, 2024

    Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks, 2024. URL https://arxiv.org/abs/2309.17453

    2024 arXiv

  53. [66]

    Characteristic neural ordinary differential equations, 2022

    Xingzi Xu, Ali Hasan, Khalil Elkhalil, Jie Ding, and Vahid Tarokh. Characteristic neural ordinary differential equations, 2022. URL https://arxiv.org/abs/2111.13207

    2022 arXiv

  54. [67]

    Focal self- attention for local-global interactions in vision transformers, 2021

    Jianwei Yang, Chunyuan Li, Pengchuan Zhang, Xiyang Dai, Bin Xiao, Lu Yuan, and Jianfeng Gao. Focal self- attention for local-global interactions in vision transformers, 2021. URL https://arxiv.org/abs/2107. 00641

    2021

  55. [68]

    Soaring from 4k to 400k: Extending llm’s context with activation beacon, 2024

    Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, and Zhicheng Dou. Soaring from 4k to 400k: Extending llm’s context with activation beacon, 2024. URL https://arxiv.org/abs/2401.03462

    2024 arXiv

  56. [69]

    tX i=1 dlt dyt λt,i(Ai θhi−1 + Bi θ ˆxi) # + dlt dyt Ct θht =

    Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanaz- eri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. Pytorch fsdp: Experiences on ...

    2023 arXiv

  57. [2013]

    URL https://arxiv.org/abs/1211.5063

    arXiv

  58. [2023]

    URL https://arxiv.org/abs/2310.01889

    arXiv

  59. [2024]

    URL https://arxiv.org/abs/2403.17297

    arXiv

Pith tools

Reviewed August 10, 2026 · model on record in the stance chip above.