TFT 与 ST-GNN：深度时序模型在 A 股选股中的真实收益边界

Fri, 23 Oct 2026 21:10:00 +0800

远思为径，码动价值。欢迎来到远思码技术实验室的最新技术实践记录。

LightGBM 那篇的结论是：多个弱因子合成后 RankIC 0.042。想要再往上，直觉上有两条路——捕捉时序结构（今天的形态和 20 天前有关）与捕捉截面结构（同行业的股票会互相传染）。

这两条正好对应 TFT 与 ST-GNN。这篇记录实验室在这两个模型上的实测，包括一个不太讨喜的结论：它们在 IC 上没有赢过 GBDT，只在组合稳定性上略有优势，而推理成本高一个数量级。

一、为什么深度模型在这类任务上不占优势

先说清楚约束条件，否则容易陷入"模型越大越好"的误区：

样本量陷阱：A 股每天约 5000 只股票 × 250 天 ≈ 125 万样本/年，听起来很多，但有效独立样本远小于此——同一天的股票高度相关，同一只股票的相邻日高度相关。
信噪比极低：日频截面收益的可预测部分只有万分之几，深度模型的参数量动辄百万级。
信号衰减快：因子有效期 5~20 个交易日，模型学到的模式可能在下个季度就失效。

因此深度模型在 A 股的正确姿势不是"替代 GBDT"，而是"提供 GBDT 看不到的那部分信息"：时序上的形态演变、行业间的溢出。

二、数据张量化：从面板到 (样本, 时间步, 特征)

TFT 需要三维输入。构造滑动窗口时要格外小心——窗口内不能混入标签期之后的数据。

# ml/deep/dataset.py
from __future__ import annotations

import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset

FEATURES = ["mom_5", "mom_20", "mom_60", "rev_5", "vol_20", "turn_20", "amihud"]


class WindowDataset(Dataset):
 """把面板数据切成 (lookback, n_features) 的时间窗，预测未来 horizon 日截面排名。

 关键约定：window 截止于 t 日（含），标签是 t+1..t+horizon 的收益。
 """

 def __init__(self, panel: pd.DataFrame, close: pd.DataFrame,
 lookback: int = 60, horizon: int = 5, min_hist: int = 250):
 self.lookback = lookback
 self.horizon = horizon
 # panel: 多级索引 (trade_date, symbol)，列为特征
 self.dates = sorted(panel.index.get_level_values(0).unique())
 self.panel = panel
 self.close = close

 xs, ys, ds = [], [], []
 for i, date in enumerate(self.dates):
 if i < min_hist or i + horizon >= len(self.dates):
 continue
 window = self.dates[i - lookback + 1: i + 1]
 block = panel.loc[window] # (lookback * n_symbols, n_feat)
 symbols = block.index.get_level_values(1).unique()
 arr = np.stack([
 block.xs(d, level=0).reindex(symbols)[FEATURES].to_numpy()
 for d in window
 ]) # (lookback, n_symbols, n_feat)
 arr = np.nan_to_num(arr, nan=0.0)
 xs.append(torch.tensor(arr.transpose(1, 0, 2), dtype=torch.float32)) # (n_symbols, lookback, n_feat)
 ds.append(date)

 c = close.loc[date, symbols]
 c_end = close.loc[self.dates[i + horizon], symbols]
 ret = (c_end / c - 1).rank(pct=True) * 2 - 1 # 截面排名标签
 ys.append(torch.tensor(ret.to_numpy(), dtype=torch.float32))

 self.x = xs
 self.y = ys
 self.date_index = ds

 def __len__(self) -> int:
 return len(self.x)

 def __getitem__(self, idx: int):
 return self.x[idx], self.y[idx]

内存是第一个坎：n_symbols × lookback × n_features 的 float32 张量，5000 × 60 × 7 ≈ 8.4 MB/天。1200 天就是 10 GB。所以在 48GB UMA 的机器上，实验室的折衷是 lookback=60、特征 7 个、只取沪深 300 + 中证 500 成分股（约 800 只），单天 1.3 MB，全量 1.5 GB，正好塞进内存。

三、TFT：可解释性才是它的卖点

pytorch-forecasting 的 TFT 把"变量选择网络（VSN）+ 多头注意力"包好了，我们只需要喂 TimeSeriesDataSet：

# ml/deep/tft_model.py
from __future__ import annotations

import lightning.pytorch as pl
from pytorch_forecasting import TemporalFusionTransformer, TimeSeriesDataSet
from pytorch_forecasting.metrics import QuantileLoss

CONFIG = dict(
 max_encoder_length=60,
 max_prediction_length=5, # 预测未来 5 日
 hidden_size=64,
 attention_head_size=4,
 dropout=0.2,
 hidden_continuous_size=32,
 learning_rate=1e-3,
 batch_size=256,
 max_epochs=30,
)


def build_dataset(df, predict: bool = False, training: TimeSeriesDataSet | None = None) -> TimeSeriesDataSet:
 return TimeSeriesDataSet(
 df,
 time_idx="time_idx",
 target="label",
 group_ids=["symbol"],
 max_encoder_length=CONFIG["max_encoder_length"],
 max_prediction_length=CONFIG["max_prediction_length"],
 time_varying_known_reals=["time_idx"], # 交易日序号，未来可知
 time_varying_unknown_reals=["mom_5", "mom_20", "mom_60",
 "rev_5", "vol_20", "turn_20", "amihud"],
 static_categoricals=["industry"],
 predict_mode=predict,
 training=training,
 allow_missing_timesteps=True,
 add_relative_time_idx=True,
 )


def build_trainer(gpus: int = 1) -> pl.Trainer:
 return pl.Trainer(
 max_epochs=CONFIG["max_epochs"],
 accelerator="gpu" if gpus else "cpu",
 precision="bf16-mixed", # ROCm 上 bf16 比 fp16 稳定
 gradient_clip_val=0.5,
 enable_checkpointing=False,
 logger=False,
 enable_progress_bar=True,
 )


def build_model(training_ds: TimeSeriesDataSet) -> TemporalFusionTransformer:
 return TemporalFusionTransformer.from_dataset(
 training_ds,
 learning_rate=CONFIG["learning_rate"],
 hidden_size=CONFIG["hidden_size"],
 attention_head_size=CONFIG["attention_head_size"],
 dropout=CONFIG["dropout"],
 hidden_continuous_size=CONFIG["hidden_continuous_size"],
 output_size=7, # 7 个分位数（默认 QuantileLoss）
 loss=QuantileLoss(),
 reduce_on_plateau_patience=4,
 )

TFT 真正的价值不在精度，而在可解释性输出：

def interpretation(model, trainer, dataloader) -> dict:
 raw = model.predict(dataloader, mode="raw", return_x=True, trainer=trainer)
 return model.interpret_output(raw.output, reduction="sum") # attention / 变量权重

interpretation 会给你三样东西：每个历史时间步的注意力权重（模型在关注 5 天前还是 40 天前）、每个变量的重要度（动量还是换手在起作用）、持续性权重（历史形态对未来预测的持续影响）。这三样东西能反过来修正因子设计，比 IC 高 0.005 有用得多。

四、ST-GNN：把行业关系建成图

截面信息传播的想法很自然：某只股票异动，同行业/上下游的股票往往跟随。用图神经网络表达就是：特征矩阵 + 邻接矩阵 → 消息传递 → 更新的特征。

# ml/deep/stgnn.py
from __future__ import annotations

import torch
import torch.nn.functional as F
from torch import nn
from torch_geometric.nn import GCNConv


def build_industry_adjacency(industry: torch.Tensor, n_keep: int = 4) -> torch.Tensor:
 """同一行业的股票两两连边，并做行归一化。industry: (n_symbols,) 的整数标签。"""
 n = industry.size(0)
 same = industry.unsqueeze(0) == industry.unsqueeze(1)
 adj = same.float()
 adj.fill_diagonal_(1.0)
 # 行业过大时每行只保留相似度最高的 n_keep 个（防止金融板块把噪声全拉进来）
 if n_keep:
 topv, topi = adj.topk(min(n_keep + 1, n), dim=1)
 keep = torch.zeros_like(adj)
 keep.scatter_(1, topi, 1.0)
 adj = adj * keep
 return adj / adj.sum(dim=1, keepdim=True).clamp_min(1e-9)


class STGNNBlock(nn.Module):
 """空间（GCN 传递截面信息）+ 时间（GRU 聚合时序）的交替堆叠。"""

 def __init__(self, n_feat: int, hidden: int = 32, edge_index: torch.Tensor | None = None):
 super().__init__()
 self.gcn = GCNConv(n_feat, hidden)
 self.gru = nn.GRU(hidden, hidden, batch_first=True, num_layers=1)
 self.head = nn.Linear(hidden, 1)
 self.edge_index = edge_index

 def forward(self, x: torch.Tensor) -> torch.Tensor:
 """x: (n_symbols, lookback, n_feat) -> (n_symbols,) 预测分数。"""
 n, t, f = x.shape
 feats = x.reshape(n * t, f)
 # 空间聚合：每行构造 block-diagonal 的边索引
 if self.edge_index is None:
 raise ValueError("edge_index 未设置")
 src, dst = self.edge_index
 ei = torch.cat([src + i * n for i in range(t)]).long(), torch.cat([dst + i * n for i in range(t)]).long()
 h = F.relu(self.gcn(feats, ei))
 h = h.view(n, t, -1)
 out, _ = self.gru(h)
 return self.head(out[:, -1, :]).squeeze(-1)

图怎么建，比模型怎么改更重要。实验室试过三种邻接：

图的构造	IC 均值	备注
不建图（GRU-only）	0.031	基线
同行业连边	0.036	稳定提升
行业 + 20 日收益相关性 top-3	0.038	提升但噪声变大
行业 + 供应链（供应商/客户）	0.033	数据稀疏，覆盖度不足

供应链图理论上最有价值，但 A 股的供应商客户关系披露覆盖率太低，最终 30% 的股票没有一条边。图的覆盖率决定了图模型的上限。

五、ROCm 上的训练细节

本地算力节点是 AMD Radeon 890M（gfx1150）+ 48GB UMA 内存，ROCm 7.2.3。要点：

# PyTorch 必须装 ROCm 构建版本
pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.2
pip install pytorch-forecasting pytorch-lightning torch-geometric

# 核显架构白名单覆盖（890M -> gfx1150）
export HSA_OVERRIDE_GFX_VERSION=11.5.0
python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"

precision="bf16-mixed" 是必需的。核显上 fp16 的 loss scaling 容易溢出，训练到第 10 个 epoch 左右突然 loss 变 NaN，我为此浪费了两天。
UMA 内存是把双刃剑：48GB 共享内存意味着 batch 可以开很大，但也意味着 CPU 内存吃紧时会直接 OOM。建议 export HSA_ENABLE_SDMA=0 减少内存拷贝抖动。
torch_geometric 的 scatter 操作在 ROCm 上比 CUDA 慢 20~30%，图很大时要控制每行邻居数。

实测训练时间（800 只股票 × 1200 天，lookback 60）：

模型	单 epoch	30 epoch	显存峰值
LightGBM（CPU）	—	约 40 秒	—
TFT（hidden 64）	95 秒	约 48 分钟	3.1 GB
ST-GNN（4 层）	60 秒	约 30 分钟	2.4 GB

六、结果：一个不性感的对照表

同一份特征、同一份样本、同一套 purged 时序切分，三个模型的表现：

模型	RankIC	ICIR	多空年化（未计成本）	多空最大回撤	推理耗时/天
LightGBM	0.042	0.31	16.4%	-12.7%	0.8 秒
TFT	0.038	0.29	17.9%	-9.4%	6.2 秒
ST-GNN	0.036	0.26	15.1%	-11.8%	4.5 秒

结论得诚实地说：

IC 上 GBDT 依然最好。深度模型没有"发现"新 alpha；
TFT 的多空组合回撤更小（-9.4% vs -12.7%），说明它学到的东西更"稳"，可能与时序注意力对趋势转折的敏感性有关；
深度模型的推理成本高 5~8 倍，对日频策略无所谓，但如果做分钟级就完全不可接受。

实验室最终的选择是：主策略用 LightGBM，TFT 的输出作为"风格因子"叠加——两者相关性只有 0.6 左右，合成的 IC 能达到 0.048，这是深度模型唯一真正值钱的地方：它与 GBDT 的误差不相关。

七、踩坑清单

窗口跨越停牌期：停牌 3 个月后复牌，窗口里 60 个交易日中有 40 个是 NaN，nan_to_num 填 0 会让模型以为"这只股票波动突然归零"。正确做法是给一个 is_suspended 掩码。
样本权重：不同日期的股票数量差异很大（2015 年牛市 vs 2018 年熊市），不加权会让模型偏向样本多的时期。实验室用 1 / 当日样本数 做权重。
随机种子敏感：深度模型的 IC 在不同种子下波动 ±0.006，报告结果必须给出多个种子的均值，否则就是在挑好看的种子。
predict_mode 与 training 参数：pytorch-forecasting 里构造验证集时必须传 training=training_dataset，否则特征缩放不一致，验证集 IC 会莫名其妙地低。
不要用全样本拟合归一化：GroupNormalizer 的 groups 参数决定了按什么分组归一化，用 symbol 会导致每只股票各自的尺度，跨股票不可比。

结语

写到这里，量化这条线上的核心环节都跑通了：数据 → 因子 → 回测 → 策略 → 机器学习 → 深度模型。

但研究做完不等于能跑。最后一篇回到工程：把这套流程变成每天早上 8 点自动生成信号并推送到手机的服务，用 systemd timer 而不是 crontab，因为失败重试、日志与依赖管理都要靠它。

远思为径，码动价值。深度模型不会给你 alpha，但它会给你更清楚的自我怀疑。

ST-GNN on 远思码技术实验室