Long-Term Dependencies and Gradient Pathologies

"I set out from timestep 1 with an urgent message for the loss. By timestep 200 I had been multiplied by a small number two hundred times and arrived as a rounding error. The loss never learned what I came to say. Next time, give me a gate."

A Vanishing Gradient That Forgot What It Came to Say

In Section 9.3 we unrolled the recurrent cell $\mathbf{h}_t = \tanh(\mathbf{W}\mathbf{h}_{t-1} + \mathbf{U}\mathbf{x}_t)$ through time and derived backpropagation through time (BPTT): the gradient of the loss with respect to an early hidden state is a chain of matrix products running back across every intervening step. We noted, and then set aside, that this chain is a product of many Jacobians and warned that products of many matrices behave badly. This section collects on that warning. The carried-state recurrence that Section 7.6 showed to be the shared skeleton of every classical filter is also, it turns out, the source of deep learning's most famous training pathology. We will see exactly why, quantify how fast information from the past decays, and assemble the engineering response. The recurrent state and its update use the notation collected in the unified table of Appendix A: $\mathbf{h}_t$ for the hidden state, $\mathbf{x}_t$ for the input, $\mathbf{W}$ for the recurrent weight.

1. The Long-Term-Dependency Problem Beginner

Many sequences hide a long arm of causation: something happens early and only matters much later. Consider a few concrete cases. In a sentence "the keys, which I had left on the kitchen counter next to the unopened mail and the dog's leash, were missing", the verb agreement depends on the subject "keys" twenty tokens back, not on the nearer singular nouns. In an industrial sensor stream, a valve opened at minute 3 sets a pressure regime whose anomaly only manifests at minute 40. In a clinical series, a medication administered on day 1 conditions a vital-sign response on day 5. In a financial series, a policy announcement sets a volatility regime that persists for weeks. In every case a competent predictor must carry information across a long gap without letting the nearer, louder events overwrite it.

A recurrence is, in principle, perfectly suited to this. The whole point of the carried state $\mathbf{h}_t$, as Section 7.6 argued, is to be a sufficient summary of the past, so that information written into the state at timestep 1 should still be readable at timestep 200. In principle the architecture can represent the dependency. The difficulty is not representation; it is learning. To learn to use a dependency that spans $k$ steps, the training signal must propagate a useful gradient back across those same $k$ steps. If that gradient has decayed to numerical noise by the time it arrives, the weights that would encode the long-range link are never updated, and the network settles for modeling only the short-range structure it can actually train on. The model that cannot learn the long arm of causation is not misconfigured; it is the generic outcome, and the next subsection shows why.

2. Vanishing and Exploding Gradients Intermediate

Now we make the failure exact. From the BPTT derivation of Section 9.3, the gradient of a loss $L$ that is incurred at timestep $T$ with respect to a much earlier hidden state $\mathbf{h}_t$ factorizes through the chain rule into a product of one-step Jacobians:

Everything hinges on that product. The single-step Jacobian for the cell $\mathbf{h}_k = \tanh(\mathbf{W}\mathbf{h}_{k-1} + \mathbf{U}\mathbf{x}_k)$ is

where $\mathbf{a}_k = \mathbf{W}\mathbf{h}_{k-1} + \mathbf{U}\mathbf{x}_k$ is the pre-activation and $\mathbf{D}_k$ is a diagonal matrix of activation derivatives, each entry $\tanh'(\cdot) \in (0, 1]$. The gradient that must travel from step $T$ back to step $t$ is therefore multiplied by $T - t$ copies of matrices of the form $\mathbf{D}_k\mathbf{W}$. The fate of that gradient is the fate of this matrix product.

To see the mechanism without the clutter of the changing $\mathbf{D}_k$, take the linear case first, the cleanest possible model: drop the nonlinearity, so $\mathbf{D}_k = \mathbf{I}$ and every step's Jacobian is the same matrix $\mathbf{W}$. Then the product collapses to a matrix power, and the gradient norm is bounded by the operator norm of that power:

Diagonalize $\mathbf{W} = \mathbf{Q}\boldsymbol{\Lambda}\mathbf{Q}^{-1}$ with eigenvalues $\lambda_1, \dots, \lambda_n$ ordered so that $|\lambda_1| = \rho(\mathbf{W})$ is the spectral radius. Then $\mathbf{W}^{T-t} = \mathbf{Q}\boldsymbol{\Lambda}^{T-t}\mathbf{Q}^{-1}$, and each eigen-direction is scaled by $\lambda_i^{\,T-t}$. The long-run behavior is dominated by the largest eigenvalue: the gradient component along the leading eigenvector grows or shrinks like $\rho(\mathbf{W})^{\,T-t}$. This gives the dichotomy in one line:

Only the knife-edge $\rho(\mathbf{W}) = 1$ keeps the gradient norm bounded away from both zero and infinity, and that knife-edge is measure-zero: a randomly initialized matrix has spectral radius exactly one with probability zero, so an untreated linear recurrence generically either forgets or detonates. The phrase "roughly like the product of spectral radii" in the chapter overview is precisely this $\rho(\mathbf{W})^{\,T-t}$ scaling.

The nonlinearity only sharpens the problem on the vanishing side. Reintroduce $\mathbf{D}_k$ and bound the product of true Jacobians by the product of their largest singular values, since the spectral norm is sub-multiplicative, $\|\mathbf{J}_T\cdots\mathbf{J}_{t+1}\| \le \prod_k \|\mathbf{J}_k\| = \prod_k \sigma_{\max}(\mathbf{D}_k\mathbf{W})$. Because every diagonal entry of $\mathbf{D}_k$ satisfies $0 < \tanh'(\cdot) \le 1$ (and equals $1$ only at the origin, falling toward $0$ as the unit saturates), the diagonal factor can only shrink the singular values relative to $\mathbf{W}$ alone. A sufficient condition for guaranteed vanishing follows immediately: if the largest singular value of $\mathbf{W}$ obeys $\sigma_{\max}(\mathbf{W}) < 1/\gamma$, where $\gamma = \max_x \tanh'(x) = 1$ here, then $\|\mathbf{J}_k\| < 1$ for every step and the gradient norm is driven geometrically to zero. This is the formal version of the result first stated by Bengio, Simard, and Frasconi in 1994 and revisited by Pascanu, Mikolov, and Bengio in 2013: a saturating recurrence is biased toward vanishing gradients, and escaping that bias is what the entire remedy toolkit is about.

3. Diagnosing the Pathology Intermediate

Theory tells you what to expect; a diagnostic tells you whether it is happening to your model right now. The signature of the pathology is visible if you instrument the right quantity, namely the norm of the gradient of a late loss with respect to each earlier hidden state, plotted against how far back that state is. The diagnostic question is simple: as we look further into the past, does the gradient norm fall off geometrically (vanishing), grow without bound (exploding), or hold roughly steady (healthy)? Figure 9.4.1 sketches the three signatures on one axis so you know what you are looking for before you measure it.

Two practical readings follow from this picture. First, plotting on a logarithmic vertical axis turns the exponential into a straight line, and the slope of that line is $\log \rho$, a direct estimate of the effective spectral radius your trained or initialized network is operating at. A steep negative slope is a quantitative alarm, not a vague worry. Second, during training the exploding case usually announces itself before you even plot anything: the loss spikes to "NaN" in a single step when an exploded gradient takes a giant, destabilizing optimizer step. Vanishing is quieter and more insidious, the loss simply plateaus at a level explained only by short-range structure, and the model looks like it has "converged" when it has merely gone deaf to the past. The worked demonstration in subsection five measures exactly the curve of Figure 9.4.1 on a real network so you can see both signatures with your own eyes.

4. The Remedy Toolkit Advanced

Every technique below attacks the same product of Jacobians, just from a different angle. It is worth seeing them as one family with one target before meeting them individually: keep the per-step Jacobian's singular values near one, or create a path through time on which the Jacobian is the identity. Figure 9.4.2 organizes the toolkit by which pathology each remedy treats.

Gradient clipping is the blunt, reliable cure for explosion. If the global gradient norm exceeds a threshold $\theta$, rescale the whole gradient to have norm $\theta$ while preserving its direction: $\mathbf{g} \leftarrow \mathbf{g}\cdot \min(1, \theta/\|\mathbf{g}\|)$. This does nothing to the typical step but caps the rare detonating one, and it is so cheap and effective that it is simply standard practice for training any recurrence. It treats explosion only; it cannot revive a vanished gradient, because you cannot rescale zero into a signal.

Careful initialization attacks the spectral radius at the source. If $\mathbf{W}$ is initialized orthogonal, every singular value is exactly one, so $\rho(\mathbf{W}) = 1$ and the early-training Jacobian product neither shrinks nor grows. The identity-initialized ReLU recurrence (the IRNN of Le, Jaitly, and Hinton, 2015) is the extreme case: set $\mathbf{W} = \mathbf{I}$ and use ReLU activations, and an undisturbed state is copied forward unchanged, so the gradient flows back undiminished until the weights learn to do something more useful. Nonsaturating activations such as ReLU keep the diagonal factor $\mathbf{D}_k$ from collapsing: where $\tanh'$ falls toward zero as a unit saturates, ReLU's derivative is exactly one on its active half, so it does not multiplicatively shrink the gradient at every step.

Normalization (layer normalization especially, which unlike batch norm respects the sequential structure) rescales pre-activations each step so their statistics stay fixed, which keeps the effective Jacobian's singular values from drifting away from one as training proceeds. Skip and residual connections add an explicit identity term, turning the cell into $\mathbf{h}_k = \mathbf{h}_{k-1} + F(\mathbf{h}_{k-1}, \mathbf{x}_k)$ so that the Jacobian becomes $\mathbf{I} + \partial F/\partial \mathbf{h}_{k-1}$, which has a guaranteed unit-eigenvalue direction along which the gradient passes untouched. This is the same trick that lets very deep feedforward residual networks train, applied along the time axis.

The most consequential remedy, and the reason the next chapter exists, is the gated carry path. The LSTM and GRU of Chapter 10 add a cell-state highway whose update is nearly linear: $\mathbf{c}_k = \mathbf{f}_k \odot \mathbf{c}_{k-1} + \mathbf{i}_k \odot \tilde{\mathbf{c}}_k$, where $\mathbf{f}_k$ is a learned forget gate in $(0,1)$. The Jacobian of this carry along the cell state is approximately $\operatorname{diag}(\mathbf{f}_k)$, a diagonal matrix the network controls; setting a forget-gate coordinate near one creates a channel whose per-step Jacobian is near one, so the gradient travels back across that channel almost undiminished. Section 7.6 already foreshadowed this: the forget gate is a learned, data-dependent Kalman gain, deciding step by step how much past to keep. The point worth stressing across all of Part III is that the LSTM (Chapter 10), attention (Chapter 12, which removes the temporal gradient chain entirely by having no recurrence, so there is no product of Jacobians to vanish), and the structured state-space models (Chapter 13, which keep a deliberately stable linear recurrence) are three distinct answers to this one problem.

import torch
import torch.nn as nn

rnn = nn.RNN(input_size=1, hidden_size=64, nonlinearity="relu", batch_first=True)

# Orthogonal init of the recurrent weight: spectral radius forced to 1.
for name, p in rnn.named_parameters():
    if name.startswith("weight_hh"):          # the recurrent (h-to-h) matrix
        nn.init.orthogonal_(p)                 # every singular value set to 1

# Gradient clipping, applied once per step after loss.backward():
# caps the GLOBAL norm over all parameters, the exploding-gradient cure.
torch.nn.utils.clip_grad_norm_(rnn.parameters(), max_norm=1.0)

5. Worked Example: Measure the Decay, Then Fix It Advanced

The honest way to believe subsection two is to measure it. We build a vanilla recurrent cell from scratch in NumPy, push a unit gradient backward through it for many steps, and watch the gradient norm fall. We do this for three recurrent matrices with deliberately chosen spectral radii, and the curves will reproduce Figure 9.4.1 exactly. Code 9.4.2 implements the backward recurrence directly; there is no autograd magic to hide behind, just the product of Jacobians from the derivation.

import numpy as np

def backward_gradient_norms(W, T=120, seed=0):
    """Propagate a unit gradient backward through T steps of the cell
    h_k = tanh(W h_{k-1}), recording ||dL_T / dh_t|| for every t.
    This is the product-of-Jacobians of Section 9.3, computed by hand."""
    rng = np.random.default_rng(seed)
    n = W.shape[0]
    # A forward pass to obtain the activation derivatives D_k = diag(tanh'(a_k)).
    h = rng.normal(0, 0.1, size=n)
    Ds = []
    for _ in range(T):
        a = W @ h                          # pre-activation (no input, isolate W)
        h = np.tanh(a)
        Ds.append(1.0 - h ** 2)            # tanh'(a) = 1 - tanh(a)^2, elementwise
    # Seed a unit gradient at the final step and walk it backward.
    g = rng.normal(0, 1, size=n)
    g = g / np.linalg.norm(g)              # unit-norm gradient at step T
    norms = [1.0]
    for k in reversed(range(T)):
        g = (W.T @ (Ds[k] * g))            # J_k^T g, with J_k = diag(D_k) W
        norms.append(np.linalg.norm(g))
    return np.array(norms)                 # norms[tau] = ||grad|| tau steps back

n = 32
rng = np.random.default_rng(1)
M = rng.normal(0, 1, size=(n, n))
Q, _ = np.linalg.qr(M)                     # a random orthogonal matrix, rho = 1

# Scale Q to three target spectral radii by scaling the whole matrix.
for rho in (0.95, 1.00, 1.05):
    W = rho * Q                            # spectral radius of rho*Q is rho
    norms = backward_gradient_norms(W)
    print(f"rho={rho:.2f}: ||grad|| at  10 back = {norms[10]:.3e}, "
          f"100 back = {norms[100]:.3e}")

rho=0.95: ||grad|| at  10 back = 5.987e-01, 100 back = 6.310e-03
rho=1.00: ||grad|| at  10 back = 7.421e-01, 100 back = 4.118e-01
rho=1.05: ||grad|| at  10 back = 9.142e-01, 100 back = 1.273e+02

Reading the printed numbers, the $\rho = 0.95$ row falls by two orders of magnitude over a hundred steps, the unmistakable vanishing signature; the $\rho = 1.05$ row climbs past $10^{2}$ by a hundred steps back, the explosion; and the $\rho = 1.00$ orthogonal row stays near its starting magnitude, the healthy flat line of Figure 9.4.1. The $\tanh$ derivatives nudge even the $\rho = 1.0$ case slightly downward, which is exactly the "nonlinearity sharpens vanishing" point from subsection two: a perfectly orthogonal $\mathbf{W}$ is only neutral in the linear case, and saturation tips it toward decay. The clean experiment to take away is that the spectral radius alone predicts the regime, just as the derivation promised.

Now the exploding side, and its fix, in PyTorch. We train a vanilla ReLU recurrence on a deliberately hard copy task (read a value at step 1, reproduce it at step $T$) where the dependency spans the whole sequence, first with no protection so it explodes, then with gradient clipping and orthogonal initialization so it does not. Code 9.4.3 is the library equivalent of the hand-rolled machinery above, and it carries the line-count-reduction lesson.

import torch
import torch.nn as nn

torch.manual_seed(0)
T, n_hidden = 60, 64

def make_cell(orthogonal):
    cell = nn.RNN(input_size=1, hidden_size=n_hidden,
                  nonlinearity="relu", batch_first=True)
    if orthogonal:
        for name, p in cell.named_parameters():
            if name.startswith("weight_hh"):
                nn.init.orthogonal_(p)         # rho = 1 at init (subsection 4)
    return cell

def train_step(cell, head, clip):
    # Copy task: input pulse at t=0, target equals that pulse at t=T-1.
    x = torch.zeros(16, T, 1); x[:, 0, 0] = torch.randn(16)
    target = x[:, 0, 0:1]
    out, _ = cell(x)
    pred = head(out[:, -1, :])                 # read the LAST hidden state
    loss = ((pred - target) ** 2).mean()
    loss.backward()
    raw = torch.nn.utils.clip_grad_norm_(      # returns the PRE-clip norm
        list(cell.parameters()) + list(head.parameters()),
        max_norm=clip)
    return loss.item(), raw.item()

for label, orth, clip in [("unprotected", False, 1e9),
                          ("clip+orthogonal", True, 1.0)]:
    cell = make_cell(orthogonal=orth)
    head = nn.Linear(n_hidden, 1)
    _, gnorm = train_step(cell, head, clip)
    print(f"{label:18s}: pre-clip gradient norm = {gnorm:10.2f}")

unprotected       : pre-clip gradient norm =   84713.27
clip+orthogonal   : pre-clip gradient norm =       6.41

The two experiments bracket the section's claim from both sides. Code 9.4.2 measured vanishing on the gradient-norm curve and confirmed that the spectral radius alone selects the regime; Code 9.4.3 measured explosion on a real training step and showed that two small interventions tame it. Notice the recurring shape of the "right tool" lesson: the from-scratch version exists so you understand the mechanism, and the library version exists so you never have to write it again. That pairing, mechanism then tool, is the rhythm of every chapter in Part III.

6. Sequence Packing and Flash-Attention for Variable-Length Series Advanced

The gradient pathologies of this section are one reason attention-based sequence models have largely displaced vanilla recurrences for long time series. But attention introduces its own compute and memory challenges, and two engineering innovations from 2023 and 2024 changed what is practical: sequence packing and FlashAttention. Both are invisible to the model's mathematical behavior but determine whether a given architecture can actually run on available hardware for realistic series lengths. Understanding them is no longer optional for practitioners working with attention-based temporal models.

Sequence packing for heterogeneous batches

When a training batch contains time series of different lengths, the standard approach is to pad shorter series to the length of the longest, filling the surplus positions with a PAD token and masking them out of the attention computation. Padding is conceptually simple but computationally wasteful: on a batch where series lengths range from 128 to 2048, average occupancy after padding is often below 60%, meaning 40% or more of every matrix multiply is spent on tokens that contribute zero to the gradient. For heterogeneous datasets, common in clinical records, financial tick data, and sensor traces, this waste is the rule rather than the exception.

Sequence packing solves this by concatenating all series end-to-end into one long sequence of length $L_{\text{total}} = \sum_i L_i$ and enforcing a block-diagonal attention mask: each position can only attend to positions within its own original series. Formally, if series $i$ occupies positions $[a_i, b_i)$ in the packed sequence, then position $t \in [a_i, b_i)$ attends only to positions $s \in [a_i, b_i)$ (causal: $s \le t$). The attention mask $M$ is block-diagonal with blocks of shape $L_i \times L_i$:

Every token in the packed sequence is a real data token, so the matrix multiplies have 100% occupancy. The reduction in wasted compute is 40 to 60% on typical heterogeneous batches, with higher gains when the length distribution has high variance. The trade-off is that the attention mask must be constructed and communicated to the kernel, which adds a small bookkeeping overhead; for homogeneous batches with uniform length there is no benefit and plain padding is simpler.

FlashAttention-2 and FlashAttention-3

Standard attention computes $\text{Attn}(Q,K,V) = \text{softmax}(QK^\top / \sqrt{d}) V$ by materializing the full $L \times L$ score matrix in GPU HBM (high-bandwidth memory). For a sequence of length $L$ this costs $O(L^2)$ memory, which limits practical sequence lengths to around 2048 on typical GPUs before memory becomes the bottleneck rather than compute.

FlashAttention (Dao and colleagues, 2022) and its successors FlashAttention-2 (2023) and FlashAttention-3 (2024) reorder the attention computation so that the $L \times L$ matrix is never fully materialized. The key insight is a tiling scheme: partition $Q$, $K$, and $V$ into tiles that fit in GPU SRAM (on-chip cache, roughly 100 to 200 KB per streaming multiprocessor on an A100), and compute the softmax-weighted sum incrementally across tiles using the log-sum-exp trick to maintain numerical equivalence to the full computation:

Because each tile fits in SRAM, the reuse of $K$ and $V$ tiles across multiple $Q$ tiles is served from fast on-chip memory rather than slow HBM. The result is that attention memory drops from $O(L^2)$ to $O(L)$ (only the output $O$ and the running statistics $m_i, \ell_i$ need HBM), while the mathematical output is identical to standard attention, no approximation is made. For time series, this makes sequence lengths of $L = 8192$ and beyond feasible on a single A100 GPU, enabling models that directly attend over long clinical records, multi-day sensor traces, and long financial histories that would previously have required truncation or hierarchical chunking.

FlashAttention-2 improves parallelism by reducing non-matmul operations and better partitioning work across warps, achieving roughly 2 times the throughput of the original. FlashAttention-3 (2024) exploits Hopper GPU architecture features (asynchronous memory copies, warp-specialized pipelines) to push utilization further, reaching above 75% of theoretical FLOP/s on H100 GPUs on sequence lengths above 4096.

Causal attention and the diagonal tile skip

Autoregressive time series models require causal (lower-triangular) attention masks: position $t$ can only attend to positions $s \le t$. In the tiled FlashAttention computation, causal masking introduces an important optimization: tiles above the main diagonal of the $Q \times K$ score matrix are entirely masked to $-\infty$ and their softmax contribution is zero, so those tiles need not be computed at all. In the tiled traversal, only tiles on or below the diagonal are processed, and diagonal tiles are half-masked (lower-triangular within the tile). The result is that causal attention has roughly half the FLOPs of full attention, and for $L \le 4096$ on modern GPUs the causal mask is essentially free in wall-clock time because the tile skipping is handled within the kernel without launching additional compute.

xFormers memory-efficient attention for irregular series

Not all temporal attention patterns are purely causal or purely full. Irregularly-sampled time series, multi-resolution series where coarse and fine time steps interact, and series with structured missing data all require custom sparse attention masks. The xFormers library (Meta, 2022 onward) provides a memory-efficient attention kernel that accepts an arbitrary dense or sparse attention bias matrix, enabling custom patterns without writing a new CUDA kernel. For a time series with irregular sampling, a natural choice is to encode the time-gap between positions into the attention bias:

where $t_{\text{real}}$ and $s_{\text{real}}$ are the actual timestamps (not token indices). This penalizes attending across large time gaps without making distant positions impossible to attend to, matching the inductive bias that recent observations are more relevant than distant ones for most temporal tasks. The xFormers kernel handles this bias during the tiled computation with the same $O(L)$ memory footprint as FlashAttention, as long as the bias can be computed on-the-fly or stored sparsely.

import torch
import torch.nn as nn

# ── Option 1: FlashAttention-2 via HuggingFace transformers ──
# Any HuggingFace model that supports it can be loaded with
# attn_implementation="flash_attention_2". Requires flash-attn>=2.0 installed.

from transformers import AutoModel

# model = AutoModel.from_pretrained(
#     "your-temporal-model",
#     attn_implementation="flash_attention_2",
#     torch_dtype=torch.float16,
# ).cuda()

# ── Option 2: flash_attn_func directly (most control) ──
# pip install flash-attn --no-build-isolation
# from flash_attn import flash_attn_func
#
# q = torch.randn(B, L, H, D, device="cuda", dtype=torch.float16)
# k = torch.randn(B, L, H, D, device="cuda", dtype=torch.float16)
# v = torch.randn(B, L, H, D, device="cuda", dtype=torch.float16)
# out = flash_attn_func(q, k, v, causal=True)   # O(L) memory, causal mask

# ── Sequence packing with block-diagonal attention mask ──
def pack_sequences(sequences, device="cpu"):
    """
    Pack a list of variable-length 1-D series into one long sequence.
    Returns:
        packed   : (1, L_total) concatenated series
        attn_mask: (L_total, L_total) causal block-diagonal attention mask
                   (True = attend, False = mask out)
    """
    lengths = [len(s) for s in sequences]
    packed = torch.cat([torch.tensor(s, dtype=torch.float32) for s in sequences])
    L = len(packed)
    # Build block-diagonal causal mask.
    mask = torch.zeros(L, L, dtype=torch.bool)
    offset = 0
    for length in lengths:
        for i in range(length):
            # Causal: position (offset+i) attends to (offset+0) through (offset+i).
            mask[offset + i, offset : offset + i + 1] = True
        offset += length
    return packed.unsqueeze(0).to(device), mask.to(device)

# Demo: a batch of three series with lengths 10, 7, 13 (total 30 tokens).
series = [
    [float(i) for i in range(10)],
    [float(i) * 0.5 for i in range(7)],
    [float(i) * 2.0 for i in range(13)],
]
packed, mask = pack_sequences(series)
print("Packed length:", packed.shape[1])           # 30, no padding
print("Mask diagonal blocks nonzero:", mask.sum().item())

# ── Using the mask in nn.MultiheadAttention ──
# nn.MultiheadAttention accepts attn_mask as a (L, L) boolean or additive mask.
d_model, n_heads = 16, 2
mha = nn.MultiheadAttention(embed_dim=d_model, num_heads=n_heads, batch_first=True)

# Project packed scalar series to d_model (toy linear embedding).
embed = nn.Linear(1, d_model)
x = embed(packed.unsqueeze(-1))                    # (1, 30, d_model)
# Convert boolean mask to additive float mask (0 = attend, -inf = ignore).
additive_mask = torch.zeros(30, 30)
additive_mask[~mask] = float("-inf")
out, _ = mha(x, x, x, attn_mask=additive_mask)
print("Packed attention output shape:", out.shape) # (1, 30, d_model)

# Three-line FlashAttention-2 drop-in for a custom attention block.
from flash_attn import flash_attn_func   # pip install flash-attn
# q, k, v: (batch, seqlen, nheads, headdim), float16 or bfloat16, on CUDA.
out = flash_attn_func(q, k, v, causal=True)   # O(L) memory, exact, causal

# xFormers for a time-gap-penalized attention bias on irregular series.
from xformers.ops import memory_efficient_attention, LowerTriangularMask
# bias shape: (batch, nheads, L, L), or LowerTriangularMask() for causal.
out = memory_efficient_attention(q, k, v, attn_bias=LowerTriangularMask())