Inside TiRex-2: bidirectional where it's safe, causal where it matters

Three kinds of variates

Real forecasting problems rarely involve a single signal. Take a power grid operator forecasting tomorrow's load. The load itself is the target. Some related signals, like measured grid frequency or spot prices, are only observed up to now: these are past covariates. Others are known in advance: the calendar, public holidays, scheduled maintenance, or a weather forecast. These are future-known covariates.

Formally, with context length \(T\) and horizon \(O\), TiRex-2 models

$$\mathcal{P}\bigl(\,\textcolor{#b0357f}{\mathbf{X}}_{T+1:T+O} \;\bigm|\; \textcolor{#b0357f}{\mathbf{X}}_{1:T},\; \textcolor{#b07a12}{\mathbf{Y}}_{1:T},\; \textcolor{#2f8a45}{\mathbf{Z}}_{1:T+O}\bigr)$$

and represents this distribution by a set of quantiles for every future time step.

Future-known covariates are often the most valuable input, but also the trickiest. To judge what a holiday on Friday means for Thursday's load, you want to know about Friday while you are at Thursday. In other words, you want to read these covariates in both directions along time.

The target must never be read that way. That constraint comes from two places:

• Training. TiRex-2 predicts the next patch at every position, not just at the end of the context. That gives a much denser learning signal, but only if position \(l\) can't peek at the patch it is supposed to predict.
• Inference. We want to stream: append a new patch, update a fixed-size state, and emit a new forecast without recomputing the history. That only works if nothing about the past depends on the future.

So the question behind the architecture is simple: how do you give future-known covariates bidirectional context without opening a back door for the target?

The big picture

Every variate, whatever its type, is padded to length \(T+O\), scaled, split into non-overlapping patches of size \(P\), and embedded by a small residual MLP shared across all variates. This gives a grid of tokens

Unobserved positions (the future of the target and of past covariates, as well as missing values) are padded. The scaler combines instance normalization with an \(\operatorname{arcsinh}\) transform, so heavy-tailed variates can't dominate the loss. Binary variates skip the transform entirely, because standardizing a sparse 0/1 signal blows its two levels far apart.

The model then stacks \(N = 12\) blocks. Each block alternates two operations:

1. a time mixer that mixes along the patch axis \(L\), separately for each variate, and
2. a variate mixer that mixes along the variate axis \(F\), separately for each patch.

At the end, a quantile head turns every target token into \(Q\) quantiles for each step of the next patch, trained with the pinball loss on all observed time steps.

Here's one block. Hover over a variate type (or tap the buttons on mobile) to trace its path:

Time mixer: one xLSTM, two directions

The time mixer is an xLSTM. It's recurrent, so it processes patches in order and carries a fixed-size state. What differs between variates is only the direction it runs in:

• Target and past covariate tokens go through a strictly forward xLSTM. The output at patch \(l\) summarizes patches \(1,\dots,l\) and nothing after.
• Future-known tokens go through the same forward pass, and through a reverse pass over \(L, L-1, \dots, l\). A linear layer fuses the two directions.

Two details matter here. First, the reverse pass is weight-tied: it uses the same parameters \(\theta\) as the forward pass, so bidirectionality costs no extra xLSTM weights. Second, only future-known covariates pay for it. Targets and past covariates, which are the streams that keep arriving at inference time, stay forward-only.

Inside, the xLSTM stack alternates mLSTM blocks (matrix memory, fully parallelizable) with sLSTM blocks (scalar memory with state tracking). The original TiRex used sLSTM only.

Variate mixer: who may listen to whom

Mixing across variates is a different problem. The number of variates changes from one dataset to the next, and there is no natural order: "temperature" isn't variate #3 in any meaningful sense. Self-attention fits naturally here. It handles any \(F\) and is permutation-equivariant within each variate type.

Two modifications make it work for TiRex-2:

Grouped attention. Several multivariate series are packed into one batch. A block-diagonal mask keeps them from attending to each other, so nothing needs to be padded to the largest \(F\) in the batch.

An asymmetric mask. Within a series, target queries may read covariate keys, but covariate queries may not read target keys:

Row = who is reading (query), column = who is being read (key).

	key x	key y	key z
query x	✓	✓	✓
query y	✕	✓	✓
query z	✕	✓	✓

It looks like a small detail. It is what makes the whole design work.

Why the mask matters: the two-hop leak

Suppose we drop the mask. Press "what if there were no mask?" in Figure 1 to see the path. Here is what goes wrong:

1. In the variate mixer, the future-known token \(z_4\) attends to the target token \(x_4\). Now \(z_4\) contains information about the target at patch 4.
2. In the next block's time mixer, the reverse xLSTM carries \(z_4\) back to \(z_3\).
3. In that block's variate mixer, \(x_3\) attends to \(z_3\) and learns about \(x_4\), which is exactly the patch it's supposed to predict.

None of these steps is suspicious on its own. Together they form a back door. A model trained to predict the next patch at every position would be rewarded for using it, and at inference time the information it learned to rely on simply isn't there.

The asymmetric mask cuts the first hop. Information can still flow backward in time, but only along the covariate channel, and the covariate channel never contains target information. Formally, one block maps the target token at patch \(l\) to a function of

$$\mathbf{H}_{x,\,1:l},\qquad \mathbf{H}_{y,\,1:l},\qquad \mathbf{H}_{z,\,1:L},$$

and by induction over the depth, the target token at patch \(l\) never depends on any target input after \(l\), in any layer. The paper gives the full proof. To our knowledge, TiRex-2 is the first time series foundation model that uses future-known covariates bidirectionally while keeping the target strictly causal in a single forward pass.

What this buys

Streaming at constant cost

Targets and past covariates, the streams that arrive at inference time, only pass through the forward xLSTM. Each new patch is a constant-cost state update, followed by a fresh forecast. A Transformer time mixer with KV caching pays \(\mathcal{O}(L)\) per new patch instead, so its latency grows with the history. We streamed TiRex-2 over up to 32 million time steps of a synthetic process. Forecast quality stayed stable even at context lengths 4,000× beyond the 8k context used in training.

The ablations agree

Removing single components on fev-bench, averaged over up to six seeds (MASE, lower is better):

The bidirectional pass matters most exactly where it should, on tasks with future-known covariates. And the model does make use of them: withholding the future-known covariates costs more than three times as much on that subset as over all tasks.

Small and efficient

In univariate mode, where the variate mixer is inactive, TiRex-2 uses 38.4M active parameters. Multivariate forecasting activates another 44.1M. For comparison, Chronos-2 uses 120M.

For reference: the full architecture

Next postMaking multivariate training data out of univariate series →