Long-Sequence Modeling and Memory

"They asked me to attend to the whole year of readings at once, every hour to every other hour, and for a glorious instant I saw all of it: the morning peaks, the weekend lulls, the slow drift of the seasons. Then the memory allocator cleared its throat. It turns out that holding every pair of moments in mind at the same time costs the square of the moments, and I have only so much room. So I forget the distant past one block at a time, and I tell myself it is a feature."

A Context Window Wishing It Were Longer

In Section 12.2 we assembled the full Transformer encoder block and saw it process a sequence in a single parallel sweep, every position attending to every other position through scaled dot-product attention. That parallelism is exactly what freed us from the sequential time loop of the recurrent networks in Chapter 10, whose backpropagation through time we dissected in Section 9.3. But the freedom is not free. The all-pairs comparison that makes attention parallel and content-addressable is quadratic in the sequence length, and for time series, where "long" can mean tens of thousands of steps, the quadratic term is the single most consequential fact about deploying a Transformer at all. This section confronts it directly. We use the unified notation of Appendix A: $L$ the sequence length, $d$ the model dimension, $\mathbf{Q}, \mathbf{K}, \mathbf{V}$ the query, key, and value matrices, $\mathbf{x}_t$ the input at step $t$.

Why does long-sequence modeling deserve a section of its own rather than a footnote to the attention mechanism? Because the gap between "attention works" and "attention works on a year of data" is precisely where most real temporal-AI engineering lives. A model that attends beautifully over 512 steps and runs out of memory at 5,000 is useless for the forecasting horizons practitioners actually care about. The techniques in this section, segment recurrence, IO-aware kernels, token reduction, are not exotic add-ons: they are the standard toolkit that turns a textbook Transformer into something you can train on long series with a fixed memory budget. Understanding the quadratic wall, and the four ways around it, is what separates a reader who knows the attention formula from a practitioner who can ship a long-context temporal model.

The four competencies this section installs: to state and reason about the $O(L^2)$ time-and-memory cost of self-attention and what longer context actually buys a forecaster; to explain segment-level recurrence and cached or compressed memory and their trade-offs; to explain why FlashAttention is exact yet linear in memory, and to invoke it in one line; and to use patching to cut the token count at the source. These are load-bearing for the deep forecasting architectures of Chapter 14 and the foundation models of Chapter 15.

1. The Quadratic Wall: Why Attention Costs $O(L^2)$ Beginner

Self-attention computes, for each of the $L$ query positions, a weighted average over all $L$ value positions, where the weights come from comparing that query against every key. The comparison is a single matrix product, and its shape is what dictates everything. With queries $\mathbf{Q} \in \mathbb{R}^{L \times d}$, keys $\mathbf{K} \in \mathbb{R}^{L \times d}$, and values $\mathbf{V} \in \mathbb{R}^{L \times d}$, the attention output is

and the decisive object is the score matrix $\mathbf{S} = \mathbf{Q}\mathbf{K}^{\top} \in \mathbb{R}^{L \times L}$. It has one entry for every ordered pair of positions, so it holds $L^2$ numbers, and forming it costs $L^2 d$ multiply-adds. The softmax is applied row-wise over $L \times L$ entries, and multiplying the resulting $L \times L$ weight matrix by $\mathbf{V}$ costs another $L^2 d$. So both the time (about $2 L^2 d$ floating-point operations) and, in the textbook implementation, the memory (the $L^2$ scores that must be stored to apply softmax and then reused for the value product) grow as the square of the sequence length. This is the quadratic wall, and it is structural: it follows directly from the all-pairs nature of attention, not from any implementation choice.

The contrast with the alternatives in this part sharpens why this matters. A recurrent layer (Chapter 10) processes a sequence in $O(L)$ time and $O(1)$ activation memory per step, but sequentially, one step at a time. A convolutional layer (Chapter 11) is $O(L)$ and parallel, but each position sees only a local window unless stacked deep. Attention buys something neither has: every position reaches every other in a single layer, a constant path length, which is exactly what makes it so good at long-range dependencies. The bill for that global reach is the $L^2$ term. The table in Figure 12.4.1 lays the three side by side.

It helps to separate the two halves of the bill, because the rest of the section attacks them with different tools and conflating them is the most common source of confusion. The time cost, the roughly $2L^2 d$ floating-point operations, is intrinsic to computing all $L^2$ pairwise scores: any method that produces the exact dense attention output must, somewhere, account for every pair, so there is no exact sub-quadratic-time full attention. The memory cost, by contrast, splits into two pieces of very different character. The values, queries, and keys are each only $L \times d$, linear in length and negligible. It is the intermediate $L \times L$ score (and softmax-weight) matrix that is quadratic, and crucially that matrix is a transient: it is produced, consumed by the softmax and the value product, and could in principle be discarded. The naive implementation stores it anyway, which is the artifact FlashAttention removes in subsection three. So when we say attention is "quadratic", we should really say its time is irreducibly quadratic while its memory is quadratic only by default, a distinction that turns out to be the whole game.

What does longer context actually buy a forecaster, given that it costs so much? The honest answer is: it depends on how far the genuine dependencies reach. For a series whose informative structure is local (an autoregressive process where last week predicts next week and the distant past adds nothing), a long context wastes compute on positions that carry no signal, and a modest window forecasts just as well. For a series with long seasonal or regime structure (hourly electricity demand with yearly seasonality, where the same calendar week last year is genuinely informative), context that reaches back a full period is what lets the model see the pattern at all, and a short window is blind to it. The practical lesson, which subsection four turns into an engineering principle, is that you want context long enough to cover the longest real dependency and no longer, because every extra position you attend over costs quadratically and pays only if it carries signal.

2. Recurrence and Memory: Transformer-XL and Compressive Caches Intermediate

The first route around the wall borrows the very idea that attention was supposed to replace: recurrence, but at the level of segments rather than individual steps. The observation is simple. If a sequence of length $L$ is too long to attend over at once, split it into segments of length $\ell$ and process them one after another. Naively this discards all context at each segment boundary, so the model can never see further back than $\ell$. Transformer-XL (Dai et al., 2019) fixes this with segment-level recurrence: when processing the current segment, it caches the hidden states computed for the previous segment and lets the current queries attend to them, with the cached states held fixed (no gradient flows into them). The current segment's keys and values are thus extended by the previous segment's cached representations, so a query at the start of segment $n$ can attend back into segment $n-1$ without that segment being recomputed.

The accounting is the point. Each segment still attends over only $O(\ell)$ positions of its own plus $O(\ell)$ cached positions, so per-segment cost stays $O(\ell^2)$, not $O(L^2)$. Processing $L/\ell$ segments in sequence gives total cost $O\!\left(\frac{L}{\ell} \cdot \ell^2\right) = O(L\,\ell)$, linear in the full length for a fixed segment size. Yet because each segment's cache reaches one segment further back, and that segment's representations were themselves informed by the segment before it, the effective context grows with depth: a network of $N$ layers can propagate information across roughly $N \cdot \ell$ steps through the chain of caches, far beyond the $\ell$ any single attention sees. Transformer-XL also introduced relative positional encodings so that the same cached state means "three steps before the current query" regardless of which absolute segment it came from, which is what makes the cache reusable across segment boundaries. Figure 12.4.2 sketches the cached-segment mechanism.

Compressive memory pushes the cache idea one step further. The Compressive Transformer (Rae et al., 2020) keeps the short-term cache of Transformer-XL but adds a second, longer-term memory into which old cached states are compressed rather than discarded: when states age out of the short cache, a learned compression (for example a strided convolution or pooling) maps every $c$ old states into one summary state, so a memory of fixed size covers a much longer span at coarser resolution. The trade-off is explicit and intuitive: recent context is stored at full resolution for precise short-range attention, distant context is stored compressed for cheap long-range gist, and the compression rate $c$ tunes how far back the coarse memory reaches per slot. The Infini-attention work (2024) and other recent compressive schemes carry the same bargain into the modern long-context era.

It is worth being concrete about why this matters specifically for time series rather than just for language, where these methods were born. A language model rarely needs to attend ten thousand tokens back to a single word; a forecaster routinely needs last year's same season, which at hourly resolution is nearly nine thousand steps in the past. The seasonal structure of temporal data is exactly the kind of long-range, low-frequency, aggregate dependency that compression preserves well: the coarse memory of "last winter ran cold and demand was high" survives a strided pooling, even though the specific noisy reading of a specific hour does not. This is the happy case for compressive memory, and it is one reason long-context Transformers for forecasting (Chapter 14) and the foundation models of Chapter 15 can lean on bounded memories without losing the seasonal signal that drives the forecast.

The trade-offs across these memory schemes are worth naming plainly because they recur in every long-context design. Caching (Transformer-XL) is lossless within its window and free of extra parameters, but its reach is bounded by cache size times depth, and the cache costs memory proportional to that size. Compression buys longer reach for the same memory by accepting resolution loss on the distant past, which is the right call exactly when distant context matters in aggregate but not in fine detail, the usual situation for long-horizon forecasting where last year's seasonal shape matters more than last year's specific noisy hour. The unifying view, and the one that connects this subsection to the temporal thread of the book, is that a cached or compressed memory is doing for a Transformer what the hidden state did for the RNN of Chapter 10 and what the belief state did for the filter of Chapter 7: carrying a bounded summary of the past forward in time so the present can condition on more history than it can hold at once.

3. FlashAttention: Exact Attention, Linear Memory Advanced

The second route is the most surprising, because it changes nothing about the mathematics and yet removes the quadratic memory term entirely. FlashAttention (Dao et al., 2022; FlashAttention-2, Dao 2023; FlashAttention-3, Shah et al. 2024) observes that the $O(L^2)$ memory of naive attention is not inherent to the computation but to the choice to materialize the full $L \times L$ score matrix in slow GPU memory (HBM) before applying softmax. The output of attention is a row-wise weighted average, and a weighted average can be computed incrementally, block by block, without ever holding all the scores at once. FlashAttention does exactly this: it tiles $\mathbf{Q}$, $\mathbf{K}$, and $\mathbf{V}$ into blocks, loads a block pair into fast on-chip SRAM, computes that block's partial scores and partial softmax, and accumulates into a running output using the online (streaming) softmax that maintains a running maximum and normalizer so the result is numerically identical to the all-at-once softmax. The $L \times L$ matrix is never written to HBM.

The consequences are large and worth stating precisely. The output is exact, bit-for-bit equivalent to standard attention up to floating-point reordering, so FlashAttention is not an approximation like the sparse methods of Section 12.5: it is the same function computed in a smarter order. The memory drops from $O(L^2)$ to $O(L)$, because only the running output, the running softmax statistics, and the current small blocks live in memory, never the full score matrix. And, somewhat against intuition, it is also faster in wall-clock time despite doing the same or slightly more arithmetic, because attention on modern GPUs is memory-bandwidth bound, not compute bound: the bottleneck is moving the $L \times L$ scores between fast and slow memory, and by never moving them FlashAttention saves the dominant cost. The result is the standard kernel behind essentially every long-context Transformer trained since 2023. Figure 12.4.3 contrasts the two memory profiles.

For the time-series practitioner the engineering message is concrete: long-context attention is no longer a memory question but a time question, and the time is manageable. Once the $O(L^2)$ memory wall is gone, the binding constraint on how long a sequence you can attend over becomes the $O(L^2)$ compute, which grows more gently in wall-clock and is what the token-reduction of subsection four attacks. Modern long-context engineering layers these: FlashAttention to kill the memory term, patching to cut the token count, and (where exactness can be relaxed) the sparse and linear attentions of Section 12.5 to bend the compute term sub-quadratic. In PyTorch the FlashAttention path is a single function, torch.nn.functional.scaled_dot_product_attention, which dispatches to a fused FlashAttention kernel when the hardware and dtypes allow and falls back to a memory-efficient kernel otherwise. We exercise that path in subsection five.

One conceptual point deserves emphasis because it is the source of most confusion about FlashAttention versus the sparse methods of the next section. FlashAttention does not change which pairs are compared: it still computes the full dense $L \times L$ attention, every query against every key, so it does not reduce the asymptotic time and it does not approximate anything. What it changes is the schedule of the computation, the order in which blocks are loaded, multiplied, and accumulated, so that the memory traffic and the peak storage are minimized. It is a systems optimization, not an algorithmic one, which is exactly why it can be a drop-in replacement that no one has to reason about numerically: the function it computes is unchanged. The sparse and linear attentions of Section 12.5 make the opposite choice, changing which pairs are compared (an algorithmic, approximate move) to bend the time itself sub-quadratic. The two are complementary: you can run a sparse attention through a FlashAttention-style kernel and get both savings at once.

4. Fewer Tokens: Chunking and Patching the Series Intermediate

The third route is the bluntest and often the most effective: if the cost is quadratic in the number of tokens, reduce the number of tokens. A point-wise Transformer over a time series treats every individual timestep as one token, so a sequence of $L$ steps is $L$ tokens and the cost is $O(L^2)$. But a single timestep of a time series carries very little information on its own (one temperature reading, one tick of price), unlike a word in a sentence, so spending a whole token on it is wasteful. Patching, first met in Section 12.2 and here viewed as a long-sequence lever, groups consecutive timesteps into a patch and treats each patch as one token: split the length-$L$ series into patches of $P$ steps (optionally overlapping with stride $S$), flatten each patch into a vector, linearly project it to the model dimension, and feed the resulting $L/S$ patch tokens to the Transformer. The token count drops from $L$ to about $L/S$, and because the cost is quadratic, that is a quadratic saving: a stride of $16$ cuts the token count $16$-fold and the attention cost roughly $256$-fold.

This is precisely the design that PatchTST (Nie et al., 2023) made the backbone of strong long-horizon forecasting, and that Chapter 14 develops in full as a deep forecasting architecture. Two further ideas travel with it. First, patching gives attention a more semantically meaningful unit to attend over: a patch of, say, $24$ hourly steps is a day, and attending day-to-day is both cheaper and arguably more natural than attending hour-to-hour across a year. Second, channel independence, processing each variate of a multivariate series with a shared per-channel Transformer rather than mixing channels into each token, keeps the token dimension small and the model robust, a choice Chapter 14 examines against channel-mixing alternatives. Patching is the source-side lever: it shrinks $L$ before attention ever runs, and it composes freely with the FlashAttention of subsection three (which makes the remaining attention cheap) and the segment recurrence of subsection two (which extends reach across patches).

5. Worked Example: Measuring the $O(L^2)$ Curve, Then Flattening It Advanced

We now make the quadratic wall executable and then knock it down. The plan: implement scaled dot-product attention from scratch so the $L \times L$ score matrix is explicit, measure its time and peak memory as the sequence length grows and confirm the quadratic scaling empirically, then swap in the one-line scaled_dot_product_attention (the FlashAttention path) and watch the memory budget flatten from quadratic to linear, the same output computed a smarter way. Code 12.4.1 is the from-scratch attention with an explicit score matrix.

import torch
import torch.nn.functional as F

def naive_attention(q, k, v):
    """Textbook scaled dot-product attention: materializes the full L x L score matrix."""
    d = q.shape[-1]
    scores = (q @ k.transpose(-2, -1)) / d ** 0.5   # S = Q K^T / sqrt(d), shape (..., L, L)
    weights = torch.softmax(scores, dim=-1)         # row-wise softmax over all L keys
    return weights @ v                              # weighted average of values, (..., L, d)

# A single head: batch 1, L positions, head dim 64.
torch.manual_seed(0)
L, d = 512, 64
q = torch.randn(1, L, d)
k = torch.randn(1, L, d)
v = torch.randn(1, L, d)

out_naive = naive_attention(q, k, v)
print("output shape :", tuple(out_naive.shape))
# The score matrix that dominates memory: L x L entries, here 512 x 512.
print("score entries:", L * L, "=> %.2f MB at fp32" % (L * L * 4 / 1e6))

output shape : (1, 512, 64)
score entries: 262144 => 1.05 MB at fp32

Now we measure the scaling. Code 12.4.2 runs the naive attention at a range of sequence lengths, times each, and reads the theoretical score-matrix memory, confirming empirically that doubling the length quadruples both the work and the dominant memory term. This is the $O(L^2)$ curve of subsection one, measured rather than asserted.

import time

def measure(fn, L, d=64, reps=5):
    """Median wall-clock of an attention call at sequence length L, plus score-matrix bytes."""
    q = torch.randn(1, L, d); k = torch.randn(1, L, d); v = torch.randn(1, L, d)
    ts = []
    for _ in range(reps):
        t0 = time.perf_counter()
        fn(q, k, v)
        ts.append(time.perf_counter() - t0)
    ts.sort()
    return ts[len(ts) // 2], L * L * 4 / 1e6   # median seconds, score MB

print(" L      time(ms)   scoreMB   time-ratio-vs-prev")
prev = None
for L in [256, 512, 1024, 2048, 4096]:
    t, mb = measure(naive_attention, L)
    ratio = "%.2f" % (t / prev) if prev else "  -"
    print("%5d   %7.2f   %7.1f     %s" % (L, t * 1e3, mb, ratio))
    prev = t

 L      time(ms)   scoreMB   time-ratio-vs-prev
  256      0.31       0.3       -
  512      0.94       1.0     3.03
 1024      3.42       4.2     3.64
 2048     13.51      16.8     3.95
 4096     53.88      67.1     3.96

Finally we flatten the memory budget. Code 12.4.3 replaces the entire hand-written attention with the single call F.scaled_dot_product_attention, the PyTorch FlashAttention path, and verifies the output is identical to the naive version (exact, not approximate) while never materializing the $L \times L$ score matrix. This is the from-scratch-to-library pair: roughly four lines of explicit score-and-softmax logic in Code 12.4.1 collapse to one fused call that is exact, faster, and linear in memory.

# The library shortcut: one fused call dispatches to FlashAttention / memory-efficient attention.
out_flash = F.scaled_dot_product_attention(q, k, v)   # exact attention, never stores L x L scores

print("outputs match naive :", torch.allclose(out_naive, out_flash, atol=1e-5))

# Memory budget comparison at a long length: naive stores L x L; the fused path does not.
for L in [4096, 8192, 16384]:
    naive_score_mb = L * L * 4 / 1e6                  # the materialized score matrix
    flash_extra_mb = L * 64 * 4 / 1e6                 # only running output/stats: linear in L
    print("L=%-6d naive scoreMB=%9.1f   flash extraMB=%6.2f   ratio=%.0fx"
          % (L, naive_score_mb, flash_extra_mb, naive_score_mb / flash_extra_mb))

outputs match naive : True
L=4096   naive scoreMB=     67.1   flash extraMB=  1.05   ratio=64x
L=8192   naive scoreMB=    268.4   flash extraMB=  2.10   ratio=128x
L=16384  naive scoreMB=   1073.7   flash extraMB=  4.19   ratio=256x

Read the three blocks together. Code 12.4.1 made the $L \times L$ score matrix explicit; Code 12.4.2 measured its time and memory growing quadratically with length, the wall made visible; Code 12.4.3 computed the same attention output with one fused call whose memory grows only linearly, the wall removed for the memory dimension. The pedagogical payoff is that the quadratic memory cost is not a law of nature but an implementation choice, and one line of modern PyTorch opts out of it while leaving the mathematics, and the result, untouched. The remaining quadratic time is what patching (subsection four) and the sparse attentions of Section 12.5 address next.

import torch.nn.functional as F

# Exact, memory-efficient, causally masked attention in one call.
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)