用 TSMixer 分离时间与特征混合，预测 ETF 收益

# TSMixer: mixing one axis at a time


# TSMixer: mixing one axis at a time

**Docker image**: `ml4t-gpu`

`04_transformers` let one position in the window depend on another through
attention, and `05_tcn` did the same through dilated convolution. Both are devices
for relating positions. A fully connected layer already relates every input to every
other, so the question TSMixer (Chen et al., 2023) asks is not whether to use one but
along which axis to apply it.

Its answer is to alternate, and never to mix both axes in a single layer. A
**time-mixing** layer applies one shared linear map across the `LOOKBACK` days, the
same map for every feature. A **feature-mixing** layer applies a small MLP across the
features, separately at each day. Stacking the two lets a representation depend on
both axes at a cost of $T^2$ mixing weights along time and $2FH$ across features,
where a single dense layer over the flattened window would cost $(TF)^2$. The model
prints its parameter count below, so the arithmetic is checkable.

**Learning objectives**:
- Read a tensor's axes well enough to say what a permutation before a `Linear` layer
  changes about which numbers get combined.
- Build the two mixing layers and say what each one can and cannot represent: the
  time-mixing map is the same for every feature, and the feature-mixing MLP is the
  same at every day.
- Say what the residual connection and the pre-normalisation are for in a stack that
  has no recurrence and no convolution to stabilise.
- Score the result against a penalised linear map on the same flattened window,
  which is the comparison that decides whether the mixing structure earned
  anything.

**Book Reference**: Chapter 13, Section 13.6 (Alternative architectures and foundation models)

**Prerequisites**: ETF features (`case_studies/etfs/`)

```python
"""Build TSMixer with alternating time and feature mixing for return prediction."""

import os

os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")

import numpy as np
import plotly.graph_objects as go
import polars as pl
import torch
import torch.nn as nn
from dl_sequences import create_sequences_multi_asset, load_dl_dataset, train_model
from ml4t.diagnostic.metrics import cross_sectional_ic_series
from plotly.subplots import make_subplots
from sklearn.linear_model import Ridge
from sklearn.preprocessing import StandardScaler

from utils.reproducibility import set_global_seeds
from utils.style import COLORS, show_plotly_with_alt
```

```python
SEED = 42
LOOKBACK = 60
D_MODEL = 32
N_LAYERS = 2
DROPOUT = 0.1
EPOCHS = 30
BATCH_SIZE = 128
LR = 1e-3
LABEL_HORIZON = 21
```

```python

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {DEVICE}")

set_global_seeds(SEED)
torch.use_deterministic_algorithms(True)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
```

## Data Loading

We use eight fixed momentum horizons from the ETF case-study pipeline.

```python
mds = load_dl_dataset("etfs")

FEATURE_COLS = [
    "ret_5d",
    "ret_10d",
    "ret_21d",
    "ret_42d",
    "ret_63d",
    "ret_126d",
    "ret_189d",
    "ret_252d",
]
TARGET_COL = mds.label_col

missing_features = sorted(set(FEATURE_COLS) - set(mds.feature_names))
if missing_features:
    raise ValueError(f"Missing required ETF momentum features: {missing_features}")

print(f"Features ({len(FEATURE_COLS)}): {FEATURE_COLS}")
print(f"Target: {TARGET_COL}")
```

## Sequence Creation and Temporal Split

```python
df = mds.dataset.drop_nulls(subset=FEATURE_COLS + [TARGET_COL])
print(f"Rows after dropping nulls: {len(df):,}")
per_date = df.group_by(mds.date_col).len().sort(mds.date_col)
print(
    f"{df[mds.date_col].min()} to {df[mds.date_col].max()}, "
    f"{df[mds.entity_cols[0]].n_unique()} funds; funds per date "
    f"{per_date['len'].min()} to {per_date['len'].max()}, median {per_date['len'].median():.0f}"
)
print(
    f"Label {TARGET_COL}: mean {df[TARGET_COL].mean():+.5f}, "
    f"standard deviation {df[TARGET_COL].std():.5f}"
)

X, y, timestamps, symbols = create_sequences_multi_asset(
    df,
    FEATURE_COLS,
    TARGET_COL,
    LOOKBACK,
    timestamp_col=mds.date_col,
    symbol_col=mds.entity_cols[0],
)
print(f"Sequences: {X.shape[0]:,}, shape: {X.shape}")

sequence_order = np.lexsort((symbols.astype(str), timestamps))
X = np.nan_to_num(X[sequence_order], nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
y = np.nan_to_num(y[sequence_order], nan=0.0).astype(np.float32)
timestamps = timestamps[sequence_order]
symbols = symbols[sequence_order]
```

### Splitting by date, with a gap for the label horizon

The split is by date, at fixed fractions of the trading days, and an example belongs
to the partition the date it carries falls in. The label is a `LABEL_HORIZON`-day
forward return, so an example dated within that many days of a boundary has an
outcome resolved by days on the far side; those examples are dropped. Input windows
may still reach back over a boundary, which is right - at decision time the model has
every past observation available.

```python
unique_dates = np.sort(np.unique(timestamps))
train_boundary_idx = int(len(unique_dates) * 0.6)
val_boundary_idx = int(len(unique_dates) * 0.8)
train_end_date = unique_dates[train_boundary_idx]
val_end_date = unique_dates[val_boundary_idx]
train_label_cutoff = unique_dates[train_boundary_idx - LABEL_HORIZON]
val_label_cutoff = unique_dates[val_boundary_idx - LABEL_HORIZON]

train_mask = timestamps < train_label_cutoff
val_mask = (timestamps >= train_end_date) & (timestamps < val_label_cutoff)
test_mask = timestamps >= val_end_date

X_train, y_train = X[train_mask], y[train_mask]
X_val, y_val = X[val_mask], y[val_mask]
X_test, y_test = X[test_mask], y[test_mask]
test_dates, test_symbols = timestamps[test_mask], symbols[test_mask]

print(f"Train: {len(X_train):,}, Val: {len(X_val):,}, Test: {len(X_test):,}")
print(
    f"Purged {LABEL_HORIZON} target dates before each boundary: "
    f"validation starts {train_end_date}, test starts {val_end_date}"
)
```

### Cross-sectional IC helper

Mean cross-sectional Spearman IC by date - the same metric used in
`01_core_architectures`, `04_transformers` and `05_tcn`, so the comparison anchors
on one per-date rank correlation across the section's architectures.

A date's IC is undefined when a model predicts the same number for every fund on it:
the predicted ranks are all tied and there is nothing to correlate. The library
returns `NaN` for such a date, and polars treats `NaN` and null as different values,
so `drop_nulls` alone leaves it in place and one of them makes the whole mean `NaN`.
Both are filtered here, and the count of dates the mean was taken over is printed
beside it, so a model that ties often is visible rather than averaged over whichever
dates happened to survive.

```python
def cross_sectional_ic_mean(y_true, y_pred, dates, syms):
    """Mean cross-sectional Spearman IC over the dates where it is defined.

    Returns the mean and the defined/total date counts. Filters both null and NaN,
    since polars `drop_nulls` leaves NaN in place.
    """
    pred_df = pl.DataFrame({"timestamp": dates, "symbol": syms, "prediction": y_pred})
    ret_df = pl.DataFrame({"timestamp": dates, "symbol": syms, "forward_return": y_true})
    ic_per_date = cross_sectional_ic_series(
        pred_df,
        ret_df,
        pred_col="prediction",
        ret_col="forward_return",
        date_col="timestamp",
        entity_col="symbol",
    )
    defined = ic_per_date.filter(pl.col("ic").is_not_null() & pl.col("ic").is_not_nan())
    mean_ic = float(defined["ic"].mean()) if defined.height else float("nan")
    return {"ic": mean_ic, "n_defined": defined.height, "n_total": ic_per_date.height}
```

## TSMixer Architecture

TSMixer alternates between two types of MLP blocks:

1. **Time-mixing**: Transposes to `(batch, features, time)` and applies the
   same temporal projection to every feature channel
2. **Feature-mixing**: Applies an MLP along the feature axis -- each timestep
   learns cross-variate interactions

Both use pre-normalization and residual connections. This is conceptually
similar to the MLP-Mixer vision architecture, adapted for time series. The
block below follows the authors' basic TSMixer implementation: one temporal
projection followed by a two-layer feature MLP.

The mixing operations can be written as:

$$\mathbf{X}' = \mathbf{X} + \sigma\bigl(W_t \cdot \operatorname{Norm}(\mathbf{X})^\top\bigr)^\top$$

for time-mixing, and similarly without the transpose for feature-mixing.

```python
class TimeMixingMLP(nn.Module):
    """Mix information across the time dimension.

    Transposes input to (batch, features, time), applies one shared temporal
    projection, then transposes back.
    """

    def __init__(self, seq_len: int, n_features: int, dropout: float):
        super().__init__()
        self.norm = nn.LayerNorm((seq_len, n_features))
        self.temporal = nn.Linear(seq_len, seq_len)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        residual = x
        x = self.norm(x)
        x = x.permute(0, 2, 1)  # (batch, features, seq_len)
        x = self.relu(self.temporal(x))
        x = x.permute(0, 2, 1)  # (batch, seq_len, features)
        return self.dropout(x) + residual
```

### Feature-Mixing MLP

Operates directly on the feature dimension at each timestep, learning
cross-variate interactions without transposing.

```python
class FeatureMixingMLP(nn.Module):
    """Mix information across the feature dimension.

    Applies MLP on the feature axis at each timestep, enabling
    cross-variate interaction learning.
    """

    def __init__(self, seq_len: int, n_features: int, hidden_dim: int, dropout: float):
        super().__init__()
        self.norm = nn.LayerNorm((seq_len, n_features))
        self.fc1 = nn.Linear(n_features, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, n_features)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        residual = x
        x = self.norm(x)
        x = self.dropout(self.relu(self.fc1(x)))
        x = self.dropout(self.fc2(x))
        return x + residual
```

### Mixer Block

Pairs one time-mixing step with one feature-mixing step, forming the
fundamental building block of TSMixer.

```python
class MixerBlock(nn.Module):
    """One block of TSMixer: time-mixing followed by feature-mixing."""

    def __init__(self, seq_len: int, n_features: int, hidden_dim: int, dropout: float):
        super().__init__()
        self.time_mix = TimeMixingMLP(seq_len, n_features, dropout)
        self.feature_mix = FeatureMixingMLP(seq_len, n_features, hidden_dim, dropout)

    def forward(self, x):
        x = self.time_mix(x)
        x = self.feature_mix(x)
        return x
```

### TSMixer Regressor

Stacks multiple mixer blocks and applies the paper's temporal forecast
projection with output length one. A small feature adapter then maps those
per-feature forecasts to the single cross-sectional return label.

```python
class TSMixerRegressor(nn.Module):
    """TSMixer blocks with a one-step temporal head and scalar feature adapter."""

    def __init__(
        self,
        seq_len: int,
        n_features: int,
        n_blocks: int = 2,
        hidden_dim: int = 32,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.blocks = nn.Sequential(
            *[MixerBlock(seq_len, n_features, hidden_dim, dropout) for _ in range(n_blocks)]
        )
        self.temporal_head = nn.Linear(seq_len, 1)
        self.feature_head = nn.Linear(n_features, 1)

    def forward(self, x):
        # x: (batch, seq_len, n_features)
        x = self.blocks(x)
        x = self.temporal_head(x.permute(0, 2, 1)).squeeze(-1)
        return self.feature_head(x).squeeze(-1)
```

```python
set_global_seeds(SEED)
model = TSMixerRegressor(
    seq_len=LOOKBACK,
    n_features=len(FEATURE_COLS),
    n_blocks=N_LAYERS,
    hidden_dim=D_MODEL,
    dropout=DROPOUT,
).to(DEVICE)

n_params = sum(p.numel() for p in model.parameters())
print(f"TSMixer parameters: {n_params:,}")
print(f"Architecture: {N_LAYERS} mixer blocks, hidden_dim={D_MODEL}")
print(f"Input: ({LOOKBACK} timesteps, {len(FEATURE_COLS)} features)")
```

## Train TSMixer

```python
print("Training TSMixer...")
history = train_model(
    model,
    X_train,
    y_train,
    X_val,
    y_val,
    EPOCHS,
    LR,
    BATCH_SIZE,
    DEVICE,
    weight_decay=0.01,
)
```

```python
epochs_axis = list(range(1, len(history["train_loss"]) + 1))
fig = go.Figure()
fig.add_trace(
    go.Scatter(
        x=epochs_axis,
        y=history["train_loss"],
        mode="lines+markers",
        name="Train",
        line_color=COLORS["blue"],
    )
)
fig.add_trace(
    go.Scatter(
        x=epochs_axis,
        y=history["val_loss"],
        mode="lines+markers",
        name="Validation",
        line_color=COLORS["amber"],
    )
)
fig.update_layout(
    title="Training and validation error per epoch",
    xaxis_title="Epoch",
    yaxis_title="Mean squared error",
    height=470,
)
show_plotly_with_alt(
    fig,
    "A line chart of mean squared error against epoch, with one line for the training "
    "set and one for the validation set. Training stops when the validation line has "
    "gone the required number of epochs without a new minimum.",
)
```

## Evaluate on Test Set

```python
model.eval()
with torch.no_grad():
    X_test_t = torch.FloatTensor(X_test).to(DEVICE)
    y_pred = model(X_test_t).cpu().numpy()

test_mse = np.mean((y_pred - y_test) ** 2)
mixer_ic = cross_sectional_ic_mean(y_test, y_pred, test_dates, test_symbols)
test_ic = mixer_ic["ic"]

print("\nTSMixer Test Results:")
print(f"  MSE: {test_mse:.6f}")
print(f"  Spearman IC: {test_ic:.4f}", end="")
print(f"  (defined on {mixer_ic['n_defined']} of {mixer_ic['n_total']} test dates)")
```

## Ridge Baseline Comparison

Flattening the 3D input to 2D and fitting Ridge regression provides a
simple linear baseline to gauge whether TSMixer's learned mixing adds value.

```python
X_train_flat = X_train.reshape(len(X_train), -1)
X_test_flat = X_test.reshape(len(X_test), -1)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train_flat)
X_test_scaled = scaler.transform(X_test_flat)

ridge = Ridge(alpha=1.0)
ridge.fit(X_train_scaled, y_train)
y_ridge_pred = ridge.predict(X_test_scaled)

ridge_mse = np.mean((y_ridge_pred - y_test) ** 2)
ridge_ic_result = cross_sectional_ic_mean(y_test, y_ridge_pred, test_dates, test_symbols)
ridge_ic = ridge_ic_result["ic"]
zero_mse = float(np.mean(y_test**2))

print("\nRidge Baseline Results:")
print(f"  MSE: {ridge_mse:.6f}")
print(f"  Spearman IC: {ridge_ic:.4f}", end="")
print(f"  (defined on {ridge_ic_result['n_defined']} of {ridge_ic_result['n_total']} test dates)")
```

## The mixer against the linear baseline

Two questions, two panels. The left asks whether the model ordered the funds usefully
on each date; the right asks whether its predicted return levels were closer than
predicting zero. A model can do better on one and worse on the other, and both are
reported because acting on a forecast uses the ordering while fitting one minimises
the squared error.

The ridge regression is the comparison that decides anything. It sees the same window
flattened into one vector and fits a penalised linear map straight to the label: $TF$
coefficients, no hidden representation, and no notion that one axis is time and the
other is features. Whatever the alternating structure is worth has to appear as a
difference from that.

```python
model_names = ["TSMixer", "Ridge"]
ic_values = [test_ic, ridge_ic]
mse_ratios = [test_mse / zero_mse, ridge_mse / zero_mse]
bar_palette = {"TSMixer": COLORS["blue"], "Ridge": COLORS["slate"]}

fig = make_subplots(
    rows=1,
    cols=2,
    subplot_titles=("Mean cross-sectional Spearman IC", "MSE relative to zero-return forecast"),
)
for model_name, ic_value, mse_ratio in zip(model_names, ic_values, mse_ratios, strict=True):
    fig.add_trace(
        go.Bar(
            x=[model_name],
            y=[ic_value],
            name=model_name,
            marker_color=bar_palette[model_name],
            text=[f"{ic_value:.3f}"],
            textposition="outside",
            showlegend=False,
        ),
        row=1,
        col=1,
    )
    fig.add_trace(
        go.Bar(
            x=[model_name],
            y=[mse_ratio],
            name=model_name,
            marker_color=bar_palette[model_name],
            text=[f"{mse_ratio:.2f}x"],
            textposition="outside",
            showlegend=False,
        ),
        row=1,
        col=2,
    )

fig.add_hline(y=0, line_color=COLORS["neutral"], row=1, col=1)
fig.add_hline(y=1, line_dash="dot", line_color=COLORS["neutral"], row=1, col=2)
fig.update_layout(
    title="TSMixer and ridge on the same test split, ranked and levelled",
    height=480,
)
fig.update_yaxes(title_text="Spearman IC", row=1, col=1)
fig.update_yaxes(title_text="MSE / zero-return MSE", row=1, col=2)
show_plotly_with_alt(
    fig,
    "Two bar panels, one bar per model. The left panel gives each model's mean "
    "cross-sectional Spearman IC against a line at zero; the right gives its test MSE "
    "as a multiple of the zero forecast's, against a dotted line at one.",
)
```

The left panel measures cross-sectional ranking, while the right panel asks
whether either fitted model improves squared error over predicting zero. These
are distinct questions. This purged single split demonstrates the architecture;
it does not establish a stable model ranking. Section 13.9 supplies the
walk-forward comparison across datasets.

## Key takeaways

1. **A permutation before a `Linear` decides which numbers get combined.** The two
   mixing layers hold the same kind of object - a dense matrix - and differ only in
   the axis it is applied along. Reading the permutation is how you know what a block
   can represent, and it is the one thing to check when adapting this architecture to
   a different tensor layout.
2. **Each mixing layer is shared along the axis it does not mix.** One temporal map
   serves all features, and one feature MLP serves all days. That sharing is what
   keeps the mixing weights at $T^2 + 2FH$ rather than the $(TF)^2$ of a dense layer
   over the flattened window, and it is also the assumption to doubt first: it says
   the same temporal pattern matters in every feature.
3. **The residual and the pre-normalisation are load-bearing.** With no recurrence
   and no convolution, a stack of dense layers over a 60-day axis has nothing else
   keeping its scale in range; every mixing layer here is wrapped in both.
4. **The head is an adaptation, not part of the paper.** TSMixer's temporal forecast
   head produces one value per feature, because the paper forecasts every channel.
   This notebook's label is one cross-sectional return, so a small feature head maps
   those per-feature values to a scalar - a modelling choice this notebook makes and
   the reader should see.

**Known limitations.** One chronological split of one ETF panel, one label horizon,
one seed, and one block count. The comparison is against one baseline, and a single
split cannot rank architectures; `12_case_study_insights` is where these families are
compared across case studies under walk-forward validation. Deterministic PyTorch
algorithms and a fixed cuBLAS workspace make repeated execution reproduce on the same
software and GPU; another environment will differ in the final decimals.

**Next**: `07_mamba_ssm` recovers a recurrent state, which both mixing and attention
gave up, and scales linearly in the sequence length - where this mixer's
`Linear(T, T)` costs $T^2$ per feature and attention costs $T^2$ outright.
![notebook output](figures/p1_1.png)
![notebook output](figures/p1_2.png)