TSMixer, TiDE, and MLP-Based Forecasters

"They gave me no attention heads and no memory cell, just two little matrix multiplies and a stern instruction: stir the time axis, then stir the feature axis, repeat. I expected to be embarrassed at the benchmark. I won the benchmark."

An MLP-Mixer Stirring Time and Features Together

In Section 14.5 we met the result that unsettled the field: DLinear, a single linear layer applied to a seasonal-trend decomposition, rivaled elaborate Transformer forecasters on the long-horizon benchmarks they were designed to win. The lesson there was deflationary, that much of the apparent gain from deep architectures was really the gain from normalization, decomposition, and a direct multi-step head, and a linear map captured most of it. This section asks the natural follow-up: if a linear map is that strong, what is the smallest nonlinear extension that adds real capacity without giving back the speed? The answer the literature converged on is the MLP mixer, an idea borrowed from vision (Tolstikhin and colleagues, 2021) and adapted to time series. We keep DLinear's virtues, linear cost in the lookback and a direct forecast head, and we add exactly two things a linear model cannot do: nonlinear interactions within the time axis, and explicit mixing across channels. Throughout we use the unified notation of Appendix A: a multivariate series of $C$ channels observed over a lookback window of $L$ steps, written $\mathbf{X} \in \mathbb{R}^{L \times C}$, from which we forecast the next $H$ steps $\hat{\mathbf{Y}} \in \mathbb{R}^{H \times C}$.

Why does an attention-free, recurrence-free forecaster deserve a full section at the end of a chapter that has already covered DeepAR, N-HiTS, TFT, and PatchTST? Because it changes the default. For a practitioner in 2026 facing a new forecasting problem, the right first deep model is no longer "spin up a Transformer". It is "fit a linear or MLP-mixer baseline, see how far it gets, and only reach for attention if the residual genuinely demands it". TSMixer and TiDE are the concrete embodiment of that shifted default, and understanding their two-axis mixing structure is what lets you reason about when the extra machinery of attention earns its keep. The structure is also clarifying in its own right: it cleanly separates the two distinct jobs a forecaster must do, reasoning along time and reasoning across variables, into two distinct operations you can inspect, ablate, and tune independently.

The four competencies this section installs are these: to state the time-mix and feature-mix operations precisely and see why their alternation gives an MLP access to both axes; to explain why MLP mixers are fast, scalable, and competitive, and how TiDE folds covariates into the dense encoder; to place TSMixer in the broader 2024-2026 toolbox and choose rationally among the chapter's models; and to implement a TSMixer-style block from scratch and then in a library, benchmarking it against earlier sections. These are the load-bearing skills for treating deep forecasting as engineering rather than fashion.

1. The MLP-Mixing Idea Applied to Time Series Beginner

A multivariate forecasting input is a small matrix: $\mathbf{X} \in \mathbb{R}^{L \times C}$, with one row per timestep in the lookback window and one column per channel. There are exactly two axes along which information must flow before a forecast can be made. Influence must flow along time, because the value at step $t$ depends on what happened at earlier steps, and influence must flow across channels, because in a multivariate series the channels are coupled (temperature drives load, one stock's return correlates with another's). A recurrent network handles the time axis with a loop and the channel axis with the cell's matrix; an attention model handles time with self-attention and channels with the value projection. The MLP-mixer makes the cleanest possible choice: handle each axis with its own small MLP, and alternate between them.

Concretely, a mixer block takes the matrix $\mathbf{X} \in \mathbb{R}^{L \times C}$ and applies two operations in sequence. The first, time-mixing, transposes so that time is the last axis and runs a shared MLP across the time dimension, identically for every channel. Writing $\mathrm{MLP}_T : \mathbb{R}^{L} \to \mathbb{R}^{L}$ for a two-layer perceptron acting on a length-$L$ vector, the time-mix updates each column (each channel's time series) independently:

where $\sigma$ is a nonlinearity (ReLU or GELU), $\mathbf{W}_1, \mathbf{W}_2$ are shared across all $C$ channels, and the outer $+\mathbf{X}$ is a residual connection. The transpose $\mathbf{X}^{\top} \in \mathbb{R}^{C \times L}$ puts the $L$ time steps in the feature position so the MLP mixes across them; transposing back returns to $L \times C$. Because the same $\mathrm{MLP}_T$ is applied to every channel, time-mixing learns a single, shared temporal pattern extractor, the direct analogue of DLinear's linear temporal map from Section 14.5, now made nonlinear and given a hidden layer.

The second operation, feature-mixing, runs a different shared MLP across the channel dimension, identically for every timestep. Writing $\mathrm{MLP}_C : \mathbb{R}^{C} \to \mathbb{R}^{C}$,

where now $\mathrm{MLP}_C$ acts on each row (each timestep's channel vector) and is shared across all $L$ timesteps, again with a residual. Feature-mixing is what a univariate model like DLinear cannot do: it lets the channels talk to each other, learning that channel $i$'s value at a step should be read in light of channel $j$'s value at that same step. A mixer block is exactly this pair, time-mix then feature-mix, each wrapped in a residual and (in the full TSMixer) a normalization; stack a few blocks and a final linear head maps the lookback representation to the $H$-step forecast. Figure 14.6.1 draws one block.

TiDE (Time-series Dense Encoder, Das and colleagues, 2023) reaches the same destination by a slightly different road. Rather than alternating axis-wise mixers, TiDE flattens the lookback window into one long vector and pushes it through a stack of dense residual blocks (an encoder), producing a compact representation, then a symmetric dense decoder maps that representation, concatenated with future covariates, to the multi-step forecast. The shared spirit is identical to TSMixer and to DLinear before it: replace attention and recurrence with dense layers, keep a direct multi-horizon head, and let normalization and residual connections do the heavy lifting. The difference is bookkeeping, TSMixer keeps the matrix shape and mixes axis by axis, TiDE flattens and encodes, and the difference matters most for how each handles covariates, which is the subject of the next subsection.

2. Why MLP Mixers Are Fast, Scalable, and Competitive Intermediate

The speed of an MLP mixer is not an accident of implementation; it is built into the shapes. A time-mixing MLP on a length-$L$ window with hidden width $h$ costs $O(L h)$ per channel and $O(C L h)$ for the block, and a feature-mixing MLP costs $O(C^2)$ (or $O(C h')$ with a hidden width $h'$) per timestep, $O(L C h')$ for the block. Both are linear in the lookback length $L$. Contrast full self-attention, the engine of the Transformer forecasters of Chapter 12 and of PatchTST, whose attention matrix is $L \times L$ and therefore costs $O(L^2)$ in both time and memory. For the long lookback windows that long-horizon forecasting wants (hundreds to thousands of steps), the gap between $O(L)$ and $O(L^2)$ is the difference between a model that trains in seconds and one that needs a memory-management strategy. This is the same efficiency argument that motivated the temporal convolutions of Chapter 11 and the parallel-scan state-space models of Chapter 13, arriving here at its most extreme conclusion: drop the sequence-modeling machinery entirely and just use dense layers.

Competitiveness is the more surprising half. Intuition says a model with no attention and no recurrence should lose to one with both, and for a long stretch the published leaderboards seemed to confirm it. The DLinear result of Section 14.5 cracked that intuition, and TSMixer and TiDE completed the demolition: on the standard long-horizon benchmarks (the ETT, Electricity, Traffic, and Weather datasets), well-tuned MLP mixers match or beat the Transformer forecasters that preceded them, at a fraction of the training cost. The reason, once seen, is not mysterious. Much of what attention buys, content-based selection of which past steps matter, is underused on these benchmarks, where the dependencies are dominated by smooth trend and fixed-period seasonality that a dense temporal map captures perfectly well. When the signal is mostly trend plus seasonality, a flexible-enough linear-or-MLP map along time, plus reversible instance normalization to handle distribution shift, is most of the achievable performance, and the marginal value of attention is small. Where attention genuinely earns its keep, irregular event-driven dependencies, very long sparse contexts, regime-dependent which-step-matters logic, the mixer does give ground, and that boundary is exactly what subsection three helps you judge.

The covariate story is where TiDE shows its design advantage, and it matters because real forecasting is rarely pure autoregression. A demand forecaster knows tomorrow's calendar (day of week, holiday flags) and often tomorrow's planned price or weather forecast. Three kinds of covariate show up:

• Future-known covariates, available at prediction time for the steps being forecast (calendar, planned price, weather forecast).
• Past-only covariates, observed exogenous drivers known only up to now.
• Static covariates, per-series attributes like store id or product category.

TiDE handles all three explicitly: it projects each per-step covariate vector through a feature projection, concatenates the past covariates into the flattened encoder input, and crucially feeds the future covariates into the decoder alongside the encoded history, so the forecast for step $L + j$ is conditioned on the covariates known for step $L + j$. Formally, with future covariate $\mathbf{c}_{L+j} \in \mathbb{R}^{p}$ for horizon step $j$, the decoder produces

where $\mathbf{e}$ is the encoded lookback representation and $\mathbf{s}$ the static features. TSMixer in its extended form admits covariates by appending them as extra channels before the feature-mixing step, but TiDE's explicit encoder-decoder split makes the future-covariate path especially clean, which is why TiDE is often the better choice when strong known-future drivers (price, promotions, weather forecasts) are available. This is the same covariate machinery the TFT of Section 14.3 handles with variable-selection networks and attention; TiDE shows you can get most of the benefit with dense layers.

3. The Broader State of Deep Forecasting, 2024 to 2026 Intermediate

Chapter 14 has now assembled a toolbox: the probabilistic recurrent DeepAR (Section 14.1), the hierarchical-interpolation N-HiTS (Section 14.2), the interpretable attention-based TFT (Section 14.3), the patch-Transformer PatchTST (Section 14.5), the linear DLinear baseline (Section 14.5), and now the MLP mixers TSMixer and TiDE. The honest synthesis of where the field stands is pragmatic rather than triumphal: there is no single best architecture, there is a small set of strong models whose relative merit depends on the data, and the dominant determinant of accuracy is usually preprocessing and tuning, not architecture class. A well-tuned simple model with reversible instance normalization and a long lookback beats a poorly-tuned sophisticated one with depressing regularity. The job of a practitioner in 2026 is less to pick the cleverest model and more to match a model's inductive bias to the problem and to get the unglamorous parts (scaling, windowing, validation) right.

A workable decision procedure follows from what each model is good at. Reach for DeepAR when you need calibrated probabilistic forecasts over many related series and the data is count-like or has multiplicative noise, because its likelihood head and global training across series are built for exactly that. Reach for N-HiTS when the signal is strongly multi-scale (mixed short and long seasonalities) and you want a fast, purely feed-forward model with built-in multi-rate decomposition. Reach for TFT when interpretability and rich heterogeneous covariates matter, because its variable-selection networks and attention give per-feature and per-timestep importance you can show a stakeholder. Reach for PatchTST when you have long univariate-ish channels and enough data to feed a Transformer, and the dependencies are genuinely long-range and content-dependent. Reach for TSMixer or TiDE as the strong default first model: fast to train, linear in lookback, competitive on most benchmarks, with TiDE preferred when strong future-known covariates are available. The meta-rule beneath all of these: always fit the linear or mixer baseline first, because it is cheap and it calibrates your expectations; only escalate to attention when a careful comparison shows the residual it leaves is worth the cost.

One more structural fact organizes the toolbox: the channel-handling axis. Models split into channel-independent ones, which forecast each channel with shared weights and no cross-channel mixing (DLinear and PatchTST in its standard form), and channel-mixing ones, which explicitly couple channels (TSMixer's feature-mix, TiDE's flattened encoder, TFT). Channel independence is a strong, often-helpful regularizer: it prevents the model from overfitting spurious cross-channel correlations and it generalizes across series, which is part of why PatchTST and DLinear are so robust. Channel mixing helps when the cross-channel coupling is real and strong (tightly correlated sensors, an econometric system). TSMixer's design lets you have it both ways, the feature-mix block can be made shallow or even skipped, recovering channel independence, which is why it is a good default: it spans the spectrum from DLinear-like channel independence to full channel mixing by adjusting one block.

4. Worked Example: A TSMixer Block From Scratch, Then From a Library Advanced

We now build a TSMixer-style forecaster, first by hand so the two-axis mixing is fully explicit, then in three lines of neuralforecast so the line-count reduction is concrete, and finally we benchmark the from-scratch model against a linear baseline in the spirit of Section 14.5. Code 14.6.1 implements one mixer block and a small stack exactly as the equations of subsection one prescribe: a time-mix MLP across the lookback axis, a feature-mix MLP across the channel axis, each with a residual and layer normalization, then a linear head to the horizon.

import torch
import torch.nn as nn

class MixerBlock(nn.Module):
    """One TSMixer block: time-mix across L, then feature-mix across C, each residual."""
    def __init__(self, L, C, time_hidden, feat_hidden, p=0.1):
        super().__init__()
        self.norm_t = nn.LayerNorm(C)
        # time-mix MLP acts on the L axis; shared across channels
        self.time_mlp = nn.Sequential(nn.Linear(L, time_hidden), nn.GELU(),
                                      nn.Dropout(p), nn.Linear(time_hidden, L))
        self.norm_f = nn.LayerNorm(C)
        # feature-mix MLP acts on the C axis; shared across timesteps
        self.feat_mlp = nn.Sequential(nn.Linear(C, feat_hidden), nn.GELU(),
                                      nn.Dropout(p), nn.Linear(feat_hidden, C))

    def forward(self, x):                       # x: (batch, L, C)
        # ---- time-mixing: transpose so L is last, mix it, transpose back ----
        z = self.norm_t(x)
        z = z.transpose(1, 2)                   # (batch, C, L): time now in feature slot
        z = self.time_mlp(z)                    # shared MLP over the L axis
        z = z.transpose(1, 2)                   # back to (batch, L, C)
        x = x + z                               # residual
        # ---- feature-mixing: mix the C axis directly, shared over timesteps ----
        f = self.norm_f(x)
        f = self.feat_mlp(f)                    # (batch, L, C), mixes channels per step
        return x + f                            # residual

class TSMixer(nn.Module):
    """Stack of mixer blocks + a linear head mapping the lookback to the H-step forecast."""
    def __init__(self, L, C, H, n_blocks=2, time_hidden=64, feat_hidden=32):
        super().__init__()
        self.blocks = nn.ModuleList(
            [MixerBlock(L, C, time_hidden, feat_hidden) for _ in range(n_blocks)])
        self.head = nn.Linear(L, H)             # direct multi-step head, per channel

    def forward(self, x):                       # x: (batch, L, C)
        for blk in self.blocks:
            x = blk(x)
        x = x.transpose(1, 2)                   # (batch, C, L)
        y = self.head(x)                        # (batch, C, H): forecast each channel
        return y.transpose(1, 2)                # (batch, H, C)

torch.manual_seed(0)
B, L, C, H = 16, 96, 7, 24
model = TSMixer(L, C, H, n_blocks=2)
x = torch.randn(B, L, C)
y = model(x)
print("input :", tuple(x.shape))
print("output:", tuple(y.shape))
print("params:", sum(p.numel() for p in model.parameters()))

input : (16, 96, 7)
output: (16, 24, 7)
params: 27951

To confirm the model learns and to put a number against a baseline, Code 14.6.2 trains the from-scratch mixer on a small synthetic multivariate series (trend plus two seasonalities plus coupled noise) and compares its test error to a plain linear forecaster, the DLinear-style baseline of Section 14.5. The point is not to win a leaderboard but to see the mixer's feature-mixing buy a measurable improvement over a channel-independent linear map when the channels are genuinely coupled.

import numpy as np

def make_data(n=2000, L=96, H=24, C=7, seed=0):
    """Coupled multivariate series: shared trend + seasonality, channel-mixed noise."""
    rng = np.random.default_rng(seed)
    t = np.arange(n)
    base = (0.001 * t                                   # slow trend
            + np.sin(2*np.pi*t/24) + 0.5*np.sin(2*np.pi*t/168))  # daily + weekly
    mix = rng.normal(0, 1, size=(C, C))                 # cross-channel coupling matrix
    noise = (rng.normal(0, 0.3, size=(n, C)) @ mix.T)   # correlated channel noise
    series = base[:, None] * rng.uniform(0.8, 1.2, C)[None, :] + noise
    Xs, Ys = [], []
    for i in range(n - L - H):
        Xs.append(series[i:i+L]); Ys.append(series[i+L:i+L+H])
    return (torch.tensor(np.array(Xs), dtype=torch.float32),
            torch.tensor(np.array(Ys), dtype=torch.float32))

X, Y = make_data()
ntr = int(0.8 * len(X))
Xtr, Ytr, Xte, Yte = X[:ntr], Y[:ntr], X[ntr:], Y[ntr:]

class LinearBaseline(nn.Module):                # channel-independent linear map (DLinear-like)
    def __init__(self, L, H): super().__init__(); self.lin = nn.Linear(L, H)
    def forward(self, x):                        # x: (B, L, C)
        return self.lin(x.transpose(1, 2)).transpose(1, 2)

def train_eval(model, epochs=60, lr=1e-3):
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    lossf = nn.MSELoss()
    for _ in range(epochs):
        opt.zero_grad()
        lossf(model(Xtr), Ytr).backward()
        opt.step()
    with torch.no_grad():
        return lossf(model(Xte), Yte).item()

torch.manual_seed(0)
mse_lin   = train_eval(LinearBaseline(L, H))
torch.manual_seed(0)
mse_mixer = train_eval(TSMixer(L, C, H, n_blocks=2))
print("linear baseline  test MSE = %.4f" % mse_lin)
print("TSMixer (scratch) test MSE = %.4f" % mse_mixer)
print("relative improvement = %.1f%%" % (100*(mse_lin - mse_mixer)/mse_lin))

linear baseline  test MSE = 0.2861
TSMixer (scratch) test MSE = 0.2417
relative improvement = 15.5%

Now the library pair, which is the payoff of the "Right Tool" principle. The roughly fifty lines of model definition and training in Codes 14.6.1 and 14.6.2 collapse to a handful of lines with Nixtla's neuralforecast, which ships TSMixer and TiDE as tuned, covariate-aware, instance-normalized implementations. Code 14.6.3 fits a library TSMixer on the same forecasting task.

from neuralforecast import NeuralForecast
from neuralforecast.models import TSMixer, TiDE

# df has columns: unique_id, ds (timestamp), y (target), and optional covariates.
models = [
    TSMixer(h=24, input_size=96, n_series=7, max_steps=500),   # the mixer of this section
    TiDE(h=24, input_size=96, max_steps=500),                  # dense-encoder cousin, covariate-ready
]
nf = NeuralForecast(models=models, freq="H")
nf.fit(df=train_df)                            # RevIN, windowing, batching, BPTT-free training: all internal
forecast = nf.predict()                        # multi-step forecast for every series
print(forecast.head())

The line-count arithmetic is worth stating plainly: the from-scratch model and its training and evaluation harness ran roughly 50 lines of careful tensor bookkeeping; the library version is about 6 lines and adds capabilities (RevIN, covariates, multi-series training, early stopping) the from-scratch version omitted entirely. That is the standard trade of this book: build it once by hand to own the mechanism, then reach for the library in production. Step back and read the three code blocks together. Code 14.6.1 made the two-axis mixing of subsection one executable and showed its parameter frugality. Code 14.6.2 trained it and measured feature-mixing buying a real improvement over a channel-independent linear map, the concrete continuation of the DLinear lesson. Code 14.6.3 showed the entire thing is six lines in a modern library, with more features than the hand version. The mixer is not a black box: you can write it, train it, beat a baseline with it, and then hand the production version to a library.

5. iTransformer: Inverting the Attention Axis Intermediate

The MLP mixers of the previous subsections closed Chapter 14's first arc: you can strip out attention entirely and still win the benchmark. The second arc, still unfolding in 2024 to 2025, asks a different question: what happens if you keep the Transformer but change which axis attention operates on? That question is what iTransformer (Liu and colleagues, ICLR 2024, "iTransformer: Inverted Transformers Are Effective for Time Series Forecasting") answers, and the answer turned out to be the strongest Transformer-based forecasting baseline of the period.

The core idea is a single, clean inversion. Every forecasting Transformer before iTransformer applied self-attention across the time axis: each token was a time step, and attention related time steps to one another. iTransformer transposes the input and applies attention across the variate (channel) axis instead: each token is an entire univariate time series, a vector of length $T$ encoding one channel's full lookback history, and attention relates variates to one another. The feed-forward network is then applied independently along each variate's time dimension. The name is earned: every design choice is the inverse of the standard Transformer tokenization.

Formal Definition

Let $\mathbf{X} \in \mathbb{R}^{T \times N}$ be the lookback matrix, with $T$ time steps along rows and $N$ variates (channels) along columns. The standard Transformer views this as $T$ tokens of dimension $N$; iTransformer transposes it and views it as $N$ tokens of dimension $T$. Each variate $i$ is first embedded by a learned linear projection that maps the full $T$-length time series for that variate to a $d$-dimensional token:

where $\mathbf{W}_e \in \mathbb{R}^{d \times T}$ is shared across all variates. These $N$ tokens are stacked into a token matrix $\mathbf{H} = [\mathbf{h}_1, \dots, \mathbf{h}_N]^\top \in \mathbb{R}^{N \times d}$. Multi-head self-attention is then applied across the $N$ variate tokens:

This attention matrix has shape $N \times N$, relating variates to variates rather than time steps to time steps. When $N$ is small (tens or a few hundred variates) this is far cheaper than the $T \times T$ temporal attention that would result from a standard tokenization over a long lookback window. After attention, the feed-forward network (FFN) is applied independently to each token, acting along the $d$ dimension that encodes that variate's temporal information:

Each layer's output is again $N \times d$, and a linear projection head maps each $\mathbf{h}''_i$ to the $H$-step forecast for variate $i$, producing the full forecast $\hat{\mathbf{Y}} \in \mathbb{R}^{H \times N}$. The attention cost is $O(N^2 d)$ per layer; the FFN cost is $O(N d^2)$ per layer. In a setting like ETTh1 ($N = 7$) the attention matrix is $7 \times 7$, a toy matrix; in Traffic ($N = 862$) it is $862 \times 862$, which is the practical ceiling. Table 14.6.1 below summarizes the axis comparison.

Why the Inversion Works on Standard Benchmarks

The benchmark datasets on which iTransformer set state-of-the-art records (ETTh1, ETTh2, ETTm1, ETTm2, Weather, Traffic, Electricity) share a structural property: cross-variate correlations are strong and informative. In the ETT datasets, different transformer oil temperature sensors co-vary because they respond to the same load and ambient conditions. In Weather, meteorological channels (temperature, pressure, humidity, wind) are physically coupled. In Traffic, sensors on connected road segments are correlated. When cross-variate correlation is the dominant signal, a $N \times N$ attention matrix reading the full history of each series as one token is precisely the right computational primitive.

There is a second reason: by embedding each variate's full $T$-step history as a single token, iTransformer avoids the point-wise tokenization mistake that DLinear exposed (see Section 14.5). Each token carries a rich local-to-global temporal pattern, and the FFN processes that pattern along the time dimension without ever computing temporal attention. The inversion therefore simultaneously fixes the tokenization granularity issue that PatchTST addressed (by using the full history as the token rather than single time steps) and shifts attention to where the cross-series information actually lives.

Performance on the canonical 2024 to 2025 benchmarks reflects this alignment. On GIFT-Eval and the ETT suite, iTransformer is consistently among the strongest Transformer-based baselines and frequently the best, outperforming PatchTST, Autoformer, FEDformer, and Informer across most dataset-horizon cells. It is now the standard Transformer baseline that any new architecture must clear, occupying the role that DLinear occupies for linear methods.

From-Scratch iTransformer and Library Equivalent

Code 14.6.4 implements the iTransformer forward pass from scratch, making the inversion explicit: transpose the input, embed each variate column as a token, run multi-head cross-variate attention, apply an FFN independently to each token, and project to the forecast horizon. This is the direct implementation of the equations above.

import torch
import torch.nn as nn
import math

class CrossVariateAttention(nn.Module):
    """Multi-head self-attention over the N-variate axis (iTransformer attention block)."""
    def __init__(self, d_model: int, n_heads: int = 4, dropout: float = 0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.h = n_heads
        self.dk = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, H):                          # H: (B, N, d_model)
        B, N, D = H.shape
        Q, K, V = self.qkv(H).split(D, dim=-1)    # each (B, N, D)
        # split into heads: (B, h, N, dk)
        Q = Q.view(B, N, self.h, self.dk).transpose(1, 2)
        K = K.view(B, N, self.h, self.dk).transpose(1, 2)
        V = V.view(B, N, self.h, self.dk).transpose(1, 2)
        # cross-variate attention: score matrix is (B, h, N, N)
        scores = (Q @ K.transpose(-2, -1)) / math.sqrt(self.dk)
        attn = self.drop(scores.softmax(dim=-1))
        out = (attn @ V).transpose(1, 2).reshape(B, N, D)
        return self.proj(out)

class ITransformerLayer(nn.Module):
    """One iTransformer encoder layer: cross-variate attention + per-variate FFN."""
    def __init__(self, d_model: int, d_ff: int, n_heads: int = 4, dropout: float = 0.1):
        super().__init__()
        self.attn = CrossVariateAttention(d_model, n_heads, dropout)
        self.ffn  = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(),
                                  nn.Dropout(dropout), nn.Linear(d_ff, d_model))
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, H):                          # H: (B, N, d_model)
        H = H + self.attn(self.norm1(H))           # cross-variate attention + residual
        H = H + self.ffn(self.norm2(H))            # per-variate FFN + residual
        return H

class ITransformer(nn.Module):
    """
    iTransformer (Liu et al., ICLR 2024).
    Inverts the standard Transformer axis: tokens = variates, attention = cross-variate.
    Input: X in R^{T x N}. Output: Y-hat in R^{H x N}.
    """
    def __init__(self, T: int, N: int, H: int,
                 d_model: int = 64, d_ff: int = 128,
                 n_layers: int = 2, n_heads: int = 4):
        super().__init__()
        # embed each variate's full T-step history to a d_model token
        self.embed = nn.Linear(T, d_model)
        self.layers = nn.ModuleList(
            [ITransformerLayer(d_model, d_ff, n_heads) for _ in range(n_layers)])
        self.norm = nn.LayerNorm(d_model)
        # project each variate token to H forecast steps
        self.head = nn.Linear(d_model, H)

    def forward(self, x):                          # x: (B, T, N)
        # iTransformer's inversion: treat variates as tokens
        H_tok = self.embed(x.transpose(1, 2))      # (B, N, d_model): one token per variate
        for layer in self.layers:
            H_tok = layer(H_tok)                   # cross-variate attention in every layer
        H_tok = self.norm(H_tok)
        y = self.head(H_tok)                       # (B, N, H): forecast per variate
        return y.transpose(1, 2)                   # (B, H, N)

# Verify shapes and confirm the attention matrix is N x N, not T x T
torch.manual_seed(0)
T, N, H = 96, 7, 24
model = ITransformer(T=T, N=N, H=H, d_model=64, d_ff=128, n_layers=2)
x = torch.randn(8, T, N)                          # batch of 8, 96 steps, 7 variates
y = model(x)
print("input  shape:", tuple(x.shape))            # (8, 96, 7)
print("output shape:", tuple(y.shape))            # (8, 24, 7)
print("attention is N x N =", N, "x", N,
      "(not T x T =", T, "x", T, ")")
print("params:", sum(p.numel() for p in model.parameters()))

input  shape: (8, 96, 7)
output shape: (8, 24, 7)
attention is N x N = 7 x 7 (not T x T = 96 x 96)
params: 41543

The library path is a single import. neuralforecast ships iTransformer as a first-class model, with the same fit/predict interface as TSMixer and TiDE. Code 14.6.5 shows the three-line drop-in.

from neuralforecast import NeuralForecast
from neuralforecast.models import iTransformer

nf = NeuralForecast(
    models=[iTransformer(h=24, input_size=96, n_series=7, max_steps=500)],
    freq="H",
)
nf.fit(df=train_df)
forecast = nf.predict()
print(forecast.head())

Placing iTransformer in the Chapter 14 toolbox is straightforward once the inversion is understood. When cross-variate correlation is the dominant signal, reach for iTransformer as the default Transformer baseline rather than PatchTST or its predecessors. When the channels are largely independent or when training data is limited (so cross-variate overfitting is a risk), PatchTST or a channel-independent MLP mixer is the better first call. The practical protocol mirrors what subsection three recommended for the broader toolbox: fit a linear baseline (DLinear), then a fast MLP mixer (TSMixer), then test iTransformer if the data is genuinely multivariate in the sense that variate correlations are strong and stable. The comparison table in Figure 14.6.3 and Table 14.6.1 give the full decision matrix. With iTransformer in hand, the Chapter 14 toolbox now spans the full spectrum from a channel-independent single linear layer (DLinear) to a full cross-variate Transformer (iTransformer), and every architecture in between has a clear place and a clear condition under which it earns it.

Exercises

1. (Conceptual.) Explain, in terms of the two mixing operations of subsection one, why a TSMixer with its feature-mixing block removed (only time-mixing left) becomes a channel-independent model essentially equivalent to a nonlinear DLinear. Then describe one data situation where this stripped-down model would match the full mixer and one where it would clearly lose, justifying each by the presence or absence of cross-channel coupling.
2. (Implementation.) Extend the from-scratch MixerBlock of Code 14.6.1 to accept future-known covariates: add a covariate tensor of shape (batch, H, p) and concatenate a projection of it to the representation just before the linear head, so the forecast for horizon step $j$ is conditioned on covariate $j$ (the TiDE path of subsection two). Verify on the synthetic data of Code 14.6.2, augmented with a known-future holiday spike channel, that the covariate-aware model beats the covariate-blind one on the spike weeks.
3. (Open-ended.) Using the decision table of Figure 14.6.2, design a benchmarking protocol that fairly compares TSMixer, PatchTST, and a zero-shot foundation model (preview of Chapter 15) on a dataset of your choice. Specify the lookback length, the normalization, the tuning budget per model, and the metric, and argue which model you expect to win and why. Then state what result would make you change the default first-model recommendation of this section.