拆解 Kronos：把 K 线当成一种语言

风影同志78324
0 阅读6分钟

拆解 Kronos：把 K 线当成一种语言

本文基于本仓库代码（model/module.py、model/kronos.py）与官方权重配置整理，所有结构和超参都能在代码里对应上。

阅读顺序建议：0（动机）→ 1（总览）→ 1.5（Transformer Q/K/V 基础） → 2（Tokenizer）→ 3（语言模型）→ 3.5（三处注意力用法）→ 4（推理）→ 5-8（对比/代码地图/实测/局限）。 如果 Q/K/V、多头、RoPE 这些概念还不熟，先读 1.5，后面会轻松很多。

0. 为什么金融时序需要「专用」基础模型

GPT 之所以能用一套架构吃掉所有 NLP 任务，是因为语言离散、低噪、语法稳定。金融 K 线恰恰相反：

  • 连续且多维：每个时间步是 open/high/low/close/volume/amount 六维向量，不是标量；
  • 高噪声：信噪比极低，同样的形态后面可能涨也可能跌；
  • 非平稳：价格量纲随时间和标的完全变化，绝对价格没有意义；
  • 分布漂移：牛熊切换后数据分布本身会变。

直接把数值回归（MSE）套在 Transformer 上，模型会退化成「预测上一根收盘价的平滑外推」。Kronos 的思路是把连续 K 线离散成 token，再用语言模型的方式建模——但为了对抗金融数据的高噪声，它在「离散化」和「生成」两个环节做了专门设计：

  1. 两级层次化量化（coarse + fine），而不是单一 VQ 码本；
  2. 生成时多次采样取平均，用蒙特卡洛降低随机性带来的方差。

1. 整体架构

flowchart LR
    A[&#34;K线窗口<br/>T × 6 (OHLCV+Amount)&#34;] --> B[&#34;z-score 归一化<br/>clip ±5&#34;]
    B --> C[&#34;KronosTokenizer<br/>Encoder&#34;]
    C --> D[&#34;BSQ 量化<br/>L2 norm → sign → 20 bits&#34;]
    D --> E[&#34;s1: 高 10 bit<br/>粗粒度码&#34;]
    D --> F[&#34;s2: 低 10 bit<br/>细粒度码&#34;]
    E --> G[&#34;Kronos Transformer<br/>decoder-only, RoPE&#34;]
    F --> G
    G --> H[&#34;自回归生成<br/>s1 → s2 → 下一个时间步&#34;]
    H --> I[&#34;KronosTokenizer<br/>Decoder&#34;]
    I --> J[&#34;反归一化<br/>得到未来 OHLCV&#34;]

两个阶段是解耦的，可以分别训练、分别微调（仓库里 finetune/train_tokenizer.py 和 finetune/train_predictor.py 就是这两步）：

阶段模块训练目标参数量（base 组合）
一KronosTokenizer重建损失 + 量化损失3,958,042
二Kronos下一 token 交叉熵24.7M / 102.3M / 499.2M

关键点：一个时间步的 6 维向量被压成一个 token（带层级），而不是 6 个 token。所以语言模型的序列长度 = 时间步数 = 上下文窗口，这让 512 的上下文等价于 512 根 K 线。

1.5 前置补课：从 Q/K/V 把 Transformer 讲清楚

这一节是「模型基础」，后面所有结构都建立在这上面。如果已经熟悉可以跳过。

1.5.1 一句话版本

注意力 = 用当前时间步的「查询（Query）」，去和历史每个时间步的「键（Key）」算相似度，再按相似度加权取它们的「值（Value）」。

它是「信息路由」：决定第 tt 个位置应该从哪些历史位置、拿多少信息。

1.5.2 Q/K/V 到底是什么

输入是一个矩阵：

X∈RB×T×dmodelX \in \mathbb{R}^{B \times T \times d_{model}}

  • BB：batch size
  • TT：序列长度（对 Kronos 就是 K 线根数）
  • dmodeld_{model}：每个时间步的向量维度

三个可学习的线性投影把 XX 变成三份：

Q=XWQ,K=XWK,V=XWV,WQ,WK,WV∈Rdmodel×dmodelQ = XW_Q,\quad K = XW_K,\quad V = XW_V,\qquad W_Q,W_K,W_V \in \mathbb{R}^{d_{model}\times d_{model}}

对应的代码（model/module.py::MultiHeadAttentionWithRoPE）：

self.q_proj = nn.Linear(d_model, d_model)
self.k_proj = nn.Linear(d_model, d_model)
self.v_proj = nn.Linear(d_model, d_model)
self.out_proj = nn.Linear(d_model, d_model)

用「图书馆」类比：

符号类比在 Kronos 里代表什么
QQ你的检索需求当前这根 K 线「想问历史什么」
KK每本书的标签历史每根 K 线「能回答什么」
VV书的内容历史每根 K 线「实际携带的信息」

1.5.3 计算流程与公式

Attention(Q,K,V)=softmax ⁣(QK⊤dk+M)V\text{Attention}(Q,K,V)=\text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V

步骤形状（Kronos-small 单条 512 上下文）含义
QK⊤QK^\top[B,H,512,512][B,H,512,512]每个位置对每个位置的「匹配分数」
/dk/\sqrt{d_k}同上缩放，dk=dmodel/H=64d_k=d_{model}/H=64
+M+M同上因果掩码，未来位置置 −∞-\infty
softmax同上每行归一化成权重，和为 1
×V\times V[B,H,512,64][B,H,512,64]按权重加权求和

为什么要除 dk\sqrt{d_k}？ 点积是 dkd_k 个数相乘再相加，方差随 dkd_k 线性增长；不做缩放时 softmax 输入过大、分布接近 one-hot，梯度几乎为 0。除以 dk\sqrt{d_k} 把方差拉回 1 附近。

因果掩码（causal mask） 就是把 {j>i}\{j>i\} 的位置设成 −∞-\infty，保证第 ii 根 K 线只能看见 ≤i\le i 的历史。代码里直接交给 PyTorch：

F.scaled_dot_product_attention(q, k, v, attn_mask=..., is_causal=True)

1.5.4 多头：为什么要切分

单头只能学一种「匹配模式」。多头把 dmodeld_{model} 切成 HH 份，每份 dk=dmodel/Hd_k = d_{model}/H，各自独立算注意力，最后拼接再过一个 WOW_O：

q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)  # [B, H, T, d_k]
...
attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, self.d_model)
return self.resid_dropout(self.out_proj(attn_output))

直觉上，不同头可以分别关注「最近几根 K 线的动量」「同时间段的日内模式」「放量突破」等不同模式。Kronos-small 用 H=8H=8、dk=64d_k=64。

1.5.5 完整形状流水线（对照代码看）

以 Kronos-small 为例：B=1, T=512, dmodel=512, H=8, dk=64B=1,\ T=512,\ d_{model}=512,\ H=8,\ d_k=64。

阶段代码张量形状
输入x[1, 512, 512]
Q/K/V 投影q_proj(x) 等[1, 512, 512]
分头.view(B,T,H,d_k).transpose(1,2)[1, 8, 512, 64]
RoPEself.rotary(q, k)仍是 [1, 8, 512, 64]
注意力分数QK⊤/dkQK^\top/\sqrt{d_k}[1, 8, 512, 512]
加权输出softmax × V[1, 8, 512, 64]
合并头.transpose(1,2).view(...)[1, 512, 512]
输出投影out_proj[1, 512, 512]

注意 [1,8,512,512] 这个张量：512×512 的注意力矩阵，内存随序列长度平方增长。这就是「上下文 512 是硬上限」的根本原因，也是 basemodel 训练时容易 OOM 的直接来源（batch 16 × 8 头 × 561² × 4B ≈ 1.6 GB，只是其中一层）。

1.5.6 RoPE：为什么只旋转 Q 和 K

绝对位置编码（sin/cos 相加）会破坏「相对位置」的平移不变性。RoPE 的做法是按位置角度旋转 Q/K 的每一对维度：

qm′=Rmqm,kn′=Rnkn⇒qm′⊤kn′=f(m−n)q'_m = R_m q_m,\quad k'_n = R_n k_n \quad\Rightarrow\quad q'^\top_m k'_n = f(m-n)

即内积只与相对距离 m−nm-n 有关。这也是为什么 V 不旋转——位置信息只用于「匹配」，不用于「内容」。

# model/module.py::RotaryPositionalEmbedding
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
...
return (q * cos) + (self._rotate_half(q) * sin),
       (k * cos) + (self._rotate_half(k) * sin)

1.5.7 别忘了「非注意力」的两块积木

TransformerBlock 里除了注意力，还有两个组件（model/module.py::TransformerBlock）：

residual = x
x = self.norm1(x)                # ① Pre-Norm
attn_out = self.self_attn(x)     # ② 注意力
x = residual + attn_out          # ③ 残差

residual = x
x = self.norm2(x)
ffn_out = self.ffn(x)            # ④ FFN
x = residual + ffn_out
组件实现作用
Pre-Norm先 RMSNorm 再进子层训练更稳，深层网络不容易梯度爆炸
RMSNormx⋅rsqrt(mean(x2)+ϵ)⋅wx \cdot \text{rsqrt}(\text{mean}(x^2)+\epsilon)\cdot w只做缩放不中心化，比 LayerNorm 省算力（无 bias）
残差x = x + sublayer(x)让梯度有「高速公路」，可堆深
SwiGLUw2(silu(w1x)⊙w3x)w_2(\text{silu}(w_1x)\odot w_3x)门控 FFN，比 ReLU/GELU 效果好，代价是 3 个权重矩阵

1.5.8 参数量验算（Kronos-small = 24,741,376）

把上面的积木数一遍，正好对上官方参数量——这是检查「有没有真正看懂结构」的最好方法：

部分计算参数量
每层注意力4×(5122+512)4 \times (512^2+512)1,050,624
每层 FFN（SwiGLU）3×512×10243 \times 512 \times 10241,572,864
每层 Norm2×5122 \times 5121,024
每层小计2,624,512
8 层20,996,096
层级嵌入（s1/s2/融合）1024×512×2+(1024×512+512)1024\times512 \times 2 + (1024\times512+512)1,573,376
时间嵌入（5 个表）(60+24+7+32+13)×512(60+24+7+32+13)\times51269,632
依赖感知层（1 层交叉注意力）4×(5122+512)+5124\times(512^2+512)+5121,051,136
最终 Norm512512
双头（s1/s2 各 1024 类）2×(512×1024+1024)2\times(512\times1024+1024)1,050,624
合计24,741,376 ✅

可以看到：FFN 的参数量（1.57M）比注意力（1.05M）还多，8 层里 FFN 占了约 60%。这也是为什么「模型越大，FFN 越宽」是常见做法（base 的 ff_dim 从 1024 提到 2048）。

2. 第一阶段：KronosTokenizer

2.1 结构

# model/kronos.py
self.embed   = nn.Linear(6, 256)                   # d_in=6 → d_model=256
self.encoder = [TransformerBlock(...)] * (n_enc_layers - 1)   # 3 层
self.quant_embed = nn.Linear(256, 20)              # → codebook 维度 = s1+s2 = 20
self.tokenizer   = BSQuantizer(s1_bits=10, s2_bits=10, ...)
self.post_quant_embed_pre = nn.Linear(10, 256)     # 只用粗码
self.post_quant_embed     = nn.Linear(20, 256)     # 用全码
self.decoder = [TransformerBlock(...)] * (n_dec_layers - 1)   # 3 层
self.head    = nn.Linear(256, 6)

TransformerBlock 是标准 pre-norm 结构，但有两个细节值得注意：

组件实现说明
归一化RMSNorm无 bias、无均值中心化，比 LayerNorm 省算力（见 1.5.7）
位置编码RotaryPositionalEmbeddingRoPE，inv_freq = 10000^(-2i/d)（见 1.5.6）
注意力F.scaled_dot_product_attention(is_causal=True)自动走 FlashAttention / 内存高效实现（见 1.5.3）
FFNw2(silu(w1(x)) * w3(x))SwiGLU，ff_dim=512

2.2 为什么用 BSQ 而不是 VQ-VAE

普通 VQ 有两个老问题：码本塌缩（大部分码字没人用）和码本利用率低（需要 EMA、dead code reset 等补丁）。Kronos 用的是 Binary Spherical Quantization（BSQ, arXiv:2406.07548）：

# model/module.py  BinarySphericalQuantizer
def quantize(self, z):
    zhat = torch.where(z > 0, 1, -1)          # 逐维取符号 → ±1
    return z + (zhat - z).detach()            # 直通估计 STE

它的本质是：

  1. 先把 z L2 归一化到单位球面；
  2. 再逐维取符号，得到 ±1 的二进制串；
  3. 码本不是可学习的表，而是超立方体的 2^D 个顶点（D=20 → 约 104.8 万个码字）。

这样：

  • 不存在「码字没人用」的塌缩问题（顶点是均匀的）；
  • 不需要存 2^20 × 20 的码本矩阵；
  • 码字之间的距离天然是汉明/球面距离。

⚠️ 但码字本身没有「数值大小」语义。 位序由学出来的旋转决定， s1 索引与价格的相关系数实测是负的（见 §2.3.4）。真正把码翻译回价格的 是 post_quant_embed_pre / post_quant_embed 两个线性层 + 解码器（还能看上下文）。

2.3 层次化：s1 粗码 + s2 细码

20 个符号位被切成两段：通道 0–9（s1，1024 个码字）+ 通道 10–19（s2，1024 个码字）。

位序：s1 是低 10 位，s2 是高 10 位。 依据是 BSQuantizer.bits_to_indices 与 KronosTokenizer.indices_to_bits 都用 2 ** arange(0, 20)，通道 0 对应权重 202^0。很多二手资料写反了。

quantized_pre = quantized[:, :, :s1_bits]          # 只要粗码（通道 0-9）
z_pre = self.post_quant_embed_pre(quantized_pre)   # 单独重建
z     = self.post_quant_embed(quantized)           # 全码重建

训练时 tokenizer 输出两路重建（z_pre 和 z），损失分别是：

L=MSE(zpre,x)⏟粗粒度+MSE(z,x)⏟全码+λbsqLbsq\mathcal{L} = \underbrace{\text{MSE}(z_{pre}, x)}_{\text{粗粒度}} + \underbrace{\text{MSE}(z, x)}_{\text{全码}} + \lambda_{bsq}\mathcal{L}_{bsq}

这个「粗-细」分层不是为了省参数，而是为了让第二阶段的语言模型可以先预测大方向、再条件补细节——金融数据里「涨还是跌」比「涨多少」更可靠，分层正好匹配这种不确定性结构。

2.3.1 完整公式推导：从原始日线到整数索引

记第 tt 根日线 xt∈R6\mathbf{x}_t\in\mathbb{R}^6（[O,H,L,C,V,A]）， d=256d=256 为 tokenizer 隐层维度，B1=B2=10B_1=B_2=10，K=B1+B2=20K=B_1+B_2=20。

① 窗口归一化 + 截断（model/kronos.py:544）

\sigma_j=\sqrt{\frac1T\sum_{t=1}^{T}(x_{t,j}-\mu_j)^2}$$ $$\tilde{x}_{t,j}=\mathrm{clip}\!\left(\frac{x_{t,j}-\mu_j}{\sigma_j+10^{-5}},\,-5,\,+5\right)$$ **② 编码器 → 20 维** $$\mathbf{h}^{(0)}_t=\mathbf{W}_{\text{emb}}\tilde{\mathbf{x}}_t+\mathbf{b}_{\text{emb}}\in\mathbb{R}^{256},\qquad \mathbf{h}^{(l)}_t=\mathrm{TransformerBlock}^{(l)}\!\left(\mathbf{h}^{(l-1)}\right)$$ $$\hat{\mathbf{z}}_t=\mathbf{W}_{\text{q}}\mathbf{h}^{(3)}_t+\mathbf{b}_{\text{q}}\in\mathbb{R}^{20}$$ 注意最后一层把 256 维压到 **20 维**——这 20 个数就是全部信息预算。实测这一步的输出范数约 $5.7\times10^4$，量级完全无约束。 **③ 球面二值化（BSQ 核心）** $$\mathbf{u}_t=\frac{\hat{\mathbf{z}}_t}{\lVert\hat{\mathbf{z}}_t\rVert_2},\qquad \lVert\mathbf{u}_t\rVert_2=1$$ $$b_{t,k}=\mathrm{sign}(u_{t,k}),\quad k=0,\dots,19$$ $$\mathbf{c}_t=\frac{1}{\sqrt{K}}\mathbf{b}_t \quad\Longrightarrow\quad \lVert\mathbf{c}_t\rVert_2=\sqrt{K\cdot\tfrac1K}=1$$ 取符号不可导，用**直通估计器（STE）**（`module.py:98-101`）： $$\mathbf{c}_t=\hat{\mathbf{z}}_t+\underbrace{\left(\tfrac{1}{\sqrt K}\mathbf{b}_t-\hat{\mathbf{z}}_t\right)}_{\text{detach 掉的部分}}$$ 前向值等于 $\frac1{\sqrt{20}}\mathbf{b}_t$，反向梯度 $\partial\mathcal{L}/\partial\hat{\mathbf{z}}_t=\partial\mathcal{L}/\partial\mathbf{c}_t$，等价于把 $\mathrm{sign}$ 当恒等映射。 **④ 拆成 s1 / s2，算整数索引** $$\underbrace{b_0,\dots,b_9}_{\text{s1 粗码}}\quad\underbrace{b_{10},\dots,b_{19}}_{\text{s2 细码}}$$ $$\boxed{\;s_1=\sum_{i=0}^{9}\frac{b_i+1}{2}\,2^{i},\qquad s_2=\sum_{j=0}^{9}\frac{b_{10+j}+1}{2}\,2^{j}\;}$$ 拼接后 $S = s_1 + 2^{10}s_2$。词表大小 $2^{B_1}+2^{B_2}=2048$，而不是 $2^{20}=1{,}048{,}576$。 实测验证（000893，90 日窗口最后 5 根）： ``` encode(half=True)[0][-5:] : [429 733 933 965 943] ← s1 encode(half=True)[1][-5:] : [ 29 655 29 29 269] ← s2 full & 1023 : [429 733 933 965 943] ✓ full >> 10 : [ 29 655 29 29 269] ✓ full == s1 + 1024*s2 ✓ ``` **⑤ 两条解码路径** $$\hat{\mathbf{y}}^{\text{(s1)}}_t=\mathrm{Head}\Big(\mathrm{Dec}^{(1..3)}\big(\mathbf{W}_{\text{pre}}\,\mathbf{c}_{t,[0:10]}+\mathbf{b}_{\text{pre}}\big)\Big)_t$$ $$\hat{\mathbf{y}}^{\text{(full)}}_t=\mathrm{Head}\Big(\mathrm{Dec}^{(1..3)}\big(\mathbf{W}_{\text{all}}\,\mathbf{c}_{t,[0:20]}+\mathbf{b}_{\text{all}}\big)\Big)_t$$ $\mathbf{W}_{\text{pre}}\in\mathbb{R}^{256\times10}$，$\mathbf{W}_{\text{all}}\in\mathbb{R}^{256\times20}$，$\mathrm{Head}\in\mathbb{R}^{6\times256}$。 关键差别是**喂给解码器的球面码只有 10 维还是 20 维**——10 维时等同于把 s2 的 10 个符号全部置零。最后反归一化回价格：$y_{t,j}=\hat y_{t,j}(\sigma_j+10^{-5})+\mu_j$。 #### 2.3.2 一个真实交易日的算例（000893，2026-11-03） ``` ① 归一化：70.05 → (70.05-84.37)/10.36 = -1.382 （窗口 μ_C=84.37, σ_C=10.36） ② encoder → quant_embed：20 维向量，范数 57348.07 ③ F.normalize 后（每维理论幅度 ±1/√20 = ±0.2236）： [ 0.2809 0.0728 0.4050 0.1822 -0.2863 0.3353 -0.3159 0.1785 0.1109 0.0404 0.1824 -0.2142 0.1987 0.2206 -0.1522 -0.2163 -0.1987 -0.2129 0.1555 -0.1799 ] ④ sign()： [ 1 1 1 1 -1 1 -1 1 1 1 | 1 -1 1 1 -1 -1 -1 -1 1 -1 ] └──────── s1 段 ────────┘ └──────── s2 段 ────────┘ ⑤ s1 = 1+2+4+8+0·16+32+0·64+128+256+512 = 943 s2 = 1+0+4+8+0+0+0+0+256+0 = 269 ⑥ 送解码器的码向量 = sign/√20 = ±0.2236（每个分量等幅） ⑦ 重建：full 路径 70.61 元（误差 0.80%），s1-only 路径 70.94 元（误差 1.27%） ``` 注意 ③ 里 $u_1=0.0728$、$u_9=0.0404$ 这种**贴近 0 的分量就是脆弱比特**—— 编码器稍微一抖就会翻号，s1 索引可能从 943 跳到 935 之类。 这解释了为什么相邻交易日的 s1 索引会剧烈跳动（`429 → 733 → 933 → 965`） 而实际收盘价只差几毛钱。 #### 2.3.3 日线实测：s1 单独 vs 全码 方法：`scripts/analyze_coarse_fine.py`，90 交易日滑窗、stride 30。 **10 位能分辨多少档？** 1024 个格子覆盖 $\pm5\sigma$（共 $10\sigma$），每格宽 $10\sigma/1024=0.00977\sigma$。折算成钱随波动放大，相对精度稳定： | 股票 | 90 日 close 1σ | 每格折合 | 相对精度 | |---|---|---|---| | 000893 亚钾国际 | 10.36 元 | 0.101 元 | 0.120% | | 000792 盐湖股份 | 16.33 元 | 0.160 元 | 0.081% | | 000408 藏格矿业 | 38.19 元 | 0.373 元 | 0.059% | | 000021 深科技 | 113.26 元 | 1.106 元 | 0.215% | | 002545 东方铁塔 | 9.76 元 | 0.095 元 | 0.112% | **s2 买到什么？** 加 10 位细码后，误差稳定下降约 1/4： | 股票 | s1 平均相对误差 | 全码平均相对误差 | 下降 | R²(s1 → full) | |---|---|---|---|---| | 000893 | 0.904% | 0.670% | 25.9% | 0.99956 → 0.99976 | | 000792 | 0.659% | 0.521% | 20.9% | 0.99883 → 0.99940 | | 000408 | 0.855% | 0.636% | 25.5% | 0.99972 → 0.99985 | | 000021 | 0.826% | 0.617% | 25.3% | 0.99887 → 0.99934 | | 002545 | 0.636% | 0.479% | 24.8% | 0.99918 → 0.99951 | > **位数不是瓶颈。** 10 位理论精度 0.12%，全码理论精度 0.0001%，但实测全码误差 > 仍有 0.5%–0.7%——差了两个数量级。真正的瓶颈是 tokenizer 那个 3.96M 参数的 > 小 Transformer：它要用 20 个数同时重建 O/H/L/C/V/A 六列。 > 所以在日线尺度上 **s1 已经很够用**，再用 s2 只挤出 0.2 个百分点。 #### 2.3.4 三个常见误解 **误解一：s1 是「价格的整数部分」，s2 是「小数部分」。** 错。三条反证： | 证据 | 实测 | |---|---| | s1 索引 vs 归一化 close 的 Spearman 秩相关 | **均值 $-0.629$**（负相关！81 个窗口里仅 24.7% 达到 \|ρ\|>0.7） | | 同一个 s1 覆盖的价格区间 | `s1=431` 出现 3 次，收盘 **73.02 ~ 93.38**（跳 20.36 元） | | 同一个完整 20 位码覆盖的价格区间 | `full=276399` 出现 6 次，收盘 **70.05 ~ 81.39**（跳 16%） | **误解二：码字用得越多说明信息越多。** 90 根日线只用到 52 个不同的 s1、28 个不同的 s2、67 个不同的完整码——词表远未饱和。这说明码字是**分布式表示**，不是查表索引。 **误解三：把位数加倍就能提高精度。** 见 §2.3.3，瓶颈在编解码器不在位宽。 > **代码里的位序不一致（不在推理路径上，但要知道）** > > | 位置 | 约定 | 影响推理？ | > |---|---|---| > | `BSQuantizer.forward → bits_to_indices` | ch0 = LSB | ✅ 主路径 | > | `KronosTokenizer.indices_to_bits` | ch0 = LSB | ✅ 主路径 | > | `BSQuantizer.basis = 2**arange(19,-1,-1)` → `codes_to_indexes` / `get_codebook_entry` | ch0 = **MSB** | ❌ 仅用于 entropy 统计 | > | `HierarchicalEmbedding.split_token`：`s2 = t & 1023` | s2 = 低位 | ❌ 只在上游传复合 tensor 时触发，实际永远传 `(s1, s2)` 元组 | > > 只要不自己调 `get_codebook_entry` 做码表可视化，就不受影响；要用的话记得先反位序。 ### 2.4 量化损失：三项组合 ```python commit_loss = beta * mean((zq.detach() - z)**2) # beta=0.05 entropy_penalty = gamma0 * persample_entropy - gamma * cb_entropy # gamma0=1.0, gamma=1.1 loss = commit_loss + zeta * entropy_penalty / inv_temperature # zeta=0.05 ``` 写成公式（$\beta=0.05,\ \gamma_0=1.0,\ \gamma=1.1,\ \zeta=0.05$）： $$\mathcal{L}_{bsq}=\underbrace{\beta\cdot\frac1T\sum_t\big\lVert\mathbf{c}_t-\hat{\mathbf{z}}_t\big\rVert_2^2}_{\text{commit}} \;+\;\zeta\big(\gamma_0 H_\text{sample}-\gamma H_\text{codebook}\big)$$ $$H_\text{sample}=\frac1T\sum_t\sum_{k=0}^{19}H_2\!\left(\sigma\!\Big(-\tfrac{4\hat z_{t,k}}{\sqrt{20}}\Big)\right), \qquad H_2(p)=-p\log p-(1-p)\log(1-p)$$ $$H_\text{codebook}=\sum_{g=1}^{5}H\big(\overline{\mathbf{p}}_g\big),\qquad \overline{\mathbf{p}}_g=\frac1T\sum_t\mathrm{softmax}\big(\mathbf{d}_{t,g}\big)$$ | 项 | 权重 | 作用 | |---|---|---| | commit loss | β=0.05 | 让 encoder 输出靠近最近的码字顶点（防止编码漂移） | | per-sample entropy | γ₀=1.0 | **最小化**单根 K 线 20 个比特的熵——要求每根样本的码尽量确定 | | codebook entropy | γ=1.1 | **最大化**整张码表的使用熵——避免所有样本挤在同一批码字上 | | zeta | 0.05 | 熵项总缩放 | 两项方向相反，$\gamma=1.1>\gamma_0=1.0$ 说明作者更担心**码本塌缩**。 > 注意 commit loss 里的 $\mathbf{c}_t$ 已经被缩放成 $\pm1/\sqrt{20}$， > 所以它实际是把 $\hat{\mathbf{z}}$ 每个分量的幅度拉向 $0.2236$（范数拉向 1）， > 而不是原始的 $\pm1$。 为了防止 `2^20` 的码本熵无法计算，代码用 `group_size=4` 的子码来近似：把 20 维按 4 维一组切成 5 组，每组码本只有 `2^4=16`，用组内 softmax 概率求熵再求和（`soft_entropy_loss`）。 ### 2.5 拆分到底省了多少参数 拆分的收益体现在**第二阶段语言模型**的词表参数上（tokenizer 用 20 维线性层，不查表）。以 Kronos-base（$d=832$）为例： | 组成部分 | 拆成 s1+s2 | 不拆（20 bit 单表） | |---|---|---| | 输入嵌入 | $2\times1024\times832 = 1{,}703{,}936$ | $1{,}048{,}576\times832 = 872{,}415{,}232$ | | 融合层 `fusion_proj` | $1664\times832+832 = 1{,}385{,}280$ | — | | 双头输出 `proj_s1/s2` | $2\times(832\times1024+1024) = 1{,}705{,}984$ | $2\times832\times1{,}048{,}576 = 1{,}744{,}830{,}464$ | | **合计** | **4,795,200**（占 base 的 4.69%） | **≈1.745B** | 省掉 **99.73%** 的词表参数。代价是每根 K 线要 2 次自回归：日线 `horizon=10` → 20 步。 在 CPU 上单步约 78 ms，一个 480 决策点的回测约 10 分钟，这个交换是划算的。 --- ## 3. 第二阶段：Kronos 语言模型 ### 3.1 输入表示 ```python self.embedding = HierarchicalEmbedding(s1_bits=10, s2_bits=10, d_model=512) self.time_emb = TemporalEmbedding(d_model=512, learn_te=True) ``` **层级嵌入**：s1、s2 各有独立的 Embedding（1024×512），拼起来后过一层线性融合： $$h_t = W_f \cdot [\text{Emb}_{s1}(s1_t);\ \text{Emb}_{s2}(s2_t)] + b$$ 注意两个 embedding 都乘了 $\sqrt{d_{model}}$ —— 这是 Transformer 里词向量与位置向量的经典缩放技巧。 > **生成顺序：s1 先、s2 后。** `Kronos.forward` 先输出 `s1_logits`，采样得到 s1 后， > 把 `emb_s1(s1)` 送进 `DependencyAwareLayer` 做一次交叉注意力，才输出 `s2_logits`。 > 所以「粗/细」不只是位宽划分，也是**因果依赖**：s2 条件于 s1（详见 §3.3）。 **时间嵌入**：`minute / hour / weekday / day / month` 五个离散特征，各自查表后**相加**。`learn_te=True` 时是可学习嵌入，否则用固定 sin/cos（`FixedEmbedding`）。 > 为什么需要时间嵌入？日线虽然 5 个特征都恒定（`minute=hour=0`），但 > `weekday`（周一效应）、`day`（月初/月末）、`month`（季末调仓）仍然携带信息； > 保留完整时段编码也让同一个分词器能处理其他频率的数据。 ### 3.2 主干 ```python x = self.embedding([s1_ids, s2_ids]) x = x + self.time_emb(stamp) x = self.token_drop(x) # token dropout, p=0.1 for layer in self.transformer: # n_layers 层 x = layer(x, key_padding_mask=padding_mask) x = self.norm(x) # RMSNorm ``` 每层都是 `pre-norm + RoPE self-attn + pre-norm + SwiGLU` 的残差堆叠（各组件的作用见 1.5.7），因果掩码（`is_causal=True`）保证只能看历史。 ### 3.3 双头 + 依赖感知层（这是最有意思的部分） 普通的两级 token 预测通常是「两个独立的分类头」。Kronos 在这里插入了一个**条件依赖层**： ```python s1_logits = self.head(x) # 先预测粗码 if use_teacher_forcing: # 训练 sibling_embed = self.embedding.emb_s1(s1_targets) else: # 推理 s1_probs = softmax(s1_logits.detach()) sample_s1_ids = multinomial(s1_probs) # 采样粗码 sibling_embed = self.embedding.emb_s1(sample_s1_ids) x2 = self.dep_layer(x, sibling_embed, padding_mask) # cross-attn: Q=s1嵌入, K/V=隐状态 s2_logits = self.head.cond_forward(x2) # 再预测细码 ``` `DependencyAwareLayer` 内部是一次 cross-attention： ```python attn_out = CrossAttn(query=sibling_embed, key=x, value=x) return self.norm(hidden + attn_out) ``` 含义是：**细码的预测必须建立在「粗码已经确定」的前提下**。训练用真值 s1 做 teacher forcing，推理用采样出的 s1 做条件——这与自回归生成中「先定框架再填细节」的直觉一致。 损失是两个头的交叉熵平均： $$\mathcal{L}_{LM} = \tfrac{1}{2}\left[\text{CE}(s1) + \text{CE}(s2)\right]$$ ### 3.4 各档模型规模 | 模型 | n_layers | d_model | n_heads | ff_dim | 参数量 | max_context | 开源 | |---|---|---|---|---|---|---|---| | Kronos-mini | — | — | — | — | 4.1M | 2048 | ✅ | | **Kronos-small** | 8 | 512 | 8 | 1024 | 24,741,376 | 512 | ✅ | | **Kronos-base** | 12 | 832 | 16 | 2048 | 102,310,592 | 512 | ✅ | | Kronos-large | 18 | 1664 | 32 | 3072 | 499,246,976 | 512 | ❌ | `small` / `base` 两行直接来自官方 `config.json`；`large` 未开源，配置取自论文 Table 1，并已用本仓库代码反算验证： ```python Kronos(n_layers=18, d_model=1664, n_heads=32, ff_dim=3072).parameters() # 499,246,976 ``` 三档全部吻合到 0.1M 精度（偏差 < 5 万，属论文四舍五入）： | 模型 | 参数公式算得 | 论文/官方标注 | |---|---|---| | small | 24,741,376 | 24.7M | | base | 102,310,592 | 102.3M | | large | 499,246,976 | 499.2M | 注意 **head_dim 守恒在 52**：`base = 832/16`、`large = 1664/32`，都是 52； `small` 是 64。也就是说 base→large 是「宽度 ×2、头数 ×2、层数 ×1.5」的等比放大， 而没有改单头维度。 Tokenizer-base：`d_in=6, d_model=256, n_heads=4, ff_dim=512, 4+4 层, s1_bits=s2_bits=10`，参数量 3,958,042。 （注意：配置里的 `n_enc_layers` / `n_dec_layers` 为 4，代码中实际堆叠 `n_enc_layers - 1` = 3 个 `TransformerBlock`，其余由输入/输出投影承担。） ### 3.4.1 三档权重是怎么训出来的 从论文（arXiv 2508.02739）能确认的部分： **① 分词器只训一次，三档共用。** 三档 LM 都挂在同一个 `Kronos-Tokenizer-base` 上，分词器先训好、然后冻结， 三档语言模型都在**同一批 token 序列**上做自回归预训练。 所以 base→large 之间**不存在各自的私有分词器**，码本语义完全一致。 **② 三档是各自从零独立预训练，不是蒸馏、也不是增量放大。** - 论文明说「随模型规模增大，各任务表现持续提升，验证了时序基础模型的 scaling law」 —— 这是独立训练才有的现象；蒸馏通常不这样表述。 - 消融实验里的 `Kronos-Parallel` 等变体也是从零训的，说明作者的流程就是「改结构 → 重训」。 - 三档的 dropout 配置完全不同（见下表），也不像是从一个 checkpoint 微调出来的。 **③ 超参随规模「反向」调整：越大越保守。** | 规模 | FFN dropout | Residual dropout | Attention dropout | Token dropout | Weight decay | |---|---|---|---|---|---| | small | 0.25 | 0.25 | 0.1 | 0.1 | 0.01 | | base | 0.20 | 0.20 | 0.0 | 0.0 | 0.05 | | large | 0.00 | 0.00 | 0.0 | 0.0 | **0.10** | 前两行能和本仓库的 `config.json` 一一对上（`pretrained/Kronos-small`、`pretrained/Kronos-base`）， 足以证明这张表可信。规律很清晰：**规模越大 → 峰值学习率越低、dropout 越少、 权重衰减越强**，用小模型的容量去补偿正则、用大模型的容量去替代正则。 优化器统一是 AdamW + cosine schedule + 线性 warmup， warmup 从峰值的 10% 起、持续 **15,000 步**（三档相同）。 **④ 数据：完全相同的语料池。** 约 121 亿条 K 线观测、96,569 个标的、45+ 家交易所、7 种频率（1 分钟到周线）， 训练数据截止到 2024 年 6 月。清洗流程会按频率设阈值，切掉： 价格跳变过大的断裂、连续无量段、长期价格停滞段、以及过短的片段。 另外对 crypto / 期货 / 外汇做了**重采样加权**（这三类在原始语料里偏少）， 避免 equities 一家独大。 **⑤ 推测（论文没写，无法验证的部分）：** - **训练步数**：论文只说 warmup 15,000 步，没给总步数。按 12B 条数据、典型 batch 规模推算，large 很可能只跑了**不足一个完整 epoch**（121 亿条相比 LLM 动辄上万亿 token 的语料算小的）。这意味着三档的「有效数据曝光量」可能相近，差距主要来自容量。 - **是否共享数据顺序 / 随机种子**：没有说明。若共享，三档的可比性更好。 - **显存与并行**：论文附录写到用 3 个 k8s pod × 8 张 RTX 4090D（共 24 卡）完成全部 训练与评测。按 499M 参数 + fp32 AdamW（参数 2GB + 优化器状态 4GB + 梯度 2GB ≈ 8GB） 推算，large 用 DDP + ZeRO/FSDP 之类分片是必然的；但具体并行策略未公开。 - **训练总时长**：未公开。24 张 4090D 训 499M 模型、12B 条数据，量级上至少要几天。 一句话总结：**shared tokenizer → 同一语料 → 三种不同宽度/深度的 decoder-only Transformer 各自从零预训练，越大越保守地调正则**。这也是为什么三档权重 可以互相替换、且都吃同一套 `Kronos-Tokenizer-base`。 --- ## 3.5 注意力在 Kronos 里的三处用法 同一套 $Q/K/V$ 机制，Kronos 用了三次，但**三处的 Q/K/V 来源完全不同**——这是理解整个模型的关键。 | 位置 | 类型 | Q 来自 | K/V 来自 | 作用 | |---|---|---|---|---| | Tokenizer encoder | **自注意力**（因果）| 当前序列隐状态 | 同一条序列 | 把 6 维 K 线压成隐向量 | | Tokenizer decoder | **自注意力**（因果）| 隐状态（重建目标）| 同一条序列 | 把码字还原成 OHLCV | | LM 主干（8 层） | **自注意力**（因果）| 当前 token 表示 | 历史 token 表示 | 学习时间序列规律 | | **依赖感知层** | **交叉注意力** | **s1 的 embedding**（1 个位置）| **主干隐状态**（上下文）| 让细码预测「带着粗码去看历史」 | ### 3.5.1 自注意力 vs 交叉注意力 ``` 自注意力（self-attention）: Q = X W_Q, K = X W_K, V = X W_V # 三者来自同一个 X → 自己内部互相看 交叉注意力（cross-attention）: Q = Y W_Q, K = X W_K, V = X W_V # Q 来自 Y，K/V 来自 X → 用 Y 去「查询」X 的信息 ``` Kronos 的 `DependencyAwareLayer` 就是第二种： ```python # model/module.py class DependencyAwareLayer(nn.Module): def forward(self, hidden_states, sibling_embed, key_padding_mask=None): attn_out = self.cross_attn( query=sibling_embed, # ← Q：s1 粗码的 embedding key=hidden_states, # ← K：主干隐状态 value=hidden_states, # ← V：主干隐状态 key_padding_mask=key_padding_mask, ) return self.norm(hidden_states + attn_out) # 残差 ``` ### 3.5.2 形状上的差异 | | 自注意力（主干） | 交叉注意力（依赖感知层） | |---|---|---| | Q 长度 | $T$（整个序列）| **1**（当前步的 s1 粗码）| | K/V 长度 | $T$ | $T$（上下文隐状态）| | 注意力矩阵 | $T \times T$ | $1 \times T$ | | 计算量 | 与 $T^2$ 成正比 | 与 $T$ 成正比，很便宜 | 推理时 `decode_s2(context, sample_pre)` 的 `context` 形状是 `[B, T, d_model]`，`sample_pre` 是 `[B, 1]`，所以这一步只算一行注意力——开销远小于主干。 ### 3.5.3 为什么非要交叉注意力，不能直接拼起来？ 因为要表达的是**条件依赖**：细码 $s_2$ 的分布应当是 $p(s_2 \mid s_1, h)$，其中 $h$ 是主干隐状态。 - 如果只是把 $s_1$ 的 embedding 加到 $h$ 上：信息是「混在一起」的，模型无法针对性地检索历史； - 用交叉注意力：$s_1$ 作为 query，可以**选择性地**从 $h$ 的不同位置取信息——例如「粗码说这是放量下跌，那就重点看历史上放量下跌后的走势」。 这也是 Kronos 与「两个独立分类头」的本质区别：**粗码不是另一个输出，而是细码的输入**。 ### 3.5.4 一句话对照 > 主干的自注意力负责「读懂历史」，依赖感知层的交叉注意力负责「在已知大方向的前提下，决定细节」。 --- ## 4. 推理：从上下文到未来 K 线 ### 4.1 预处理：逐窗口 z-score ```python x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0) x = (x - x_mean) / (x_std + 1e-5) x = np.clip(x, -self.clip, self.clip) # clip=5 ``` 注意是**对输入窗口内每个特征分别标准化**，不是全局统计量——这样才能适应不同价格量级和波动率。生成结束后再反变换： $$\hat{x} = \hat{x}_{norm} \cdot (\sigma + 10^{-5}) + \mu$$ ### 4.2 自回归生成循环 ```python # model/kronos.py auto_regressive_inference x_token = tokenizer.encode(x, half=True) # → [s1_ids, s2_ids] for i in range(pred_len): s1_logits, context = model.decode_s1(pre_buffer, post_buffer, current_stamp) sample_pre = sample_from_logits(s1_logits[:, -1, :], T, top_k, top_p) s2_logits = model.decode_s2(context, sample_pre) sample_post = sample_from_logits(s2_logits[:, -1, :], T, top_k, top_p) # 写入滑动缓冲区，超出 max_context 时 roll 左移 ``` 几个工程细节： 1. **滑动窗口**：`pre_buffer/post_buffer` 固定长度 `max_context`，超出后 `torch.roll(-1)` 丢弃最老的 token； 2. **每个时间步两次采样**：先 s1 再 s2，两次都走 `temperature → top-k/top-p → multinomial`； 3. **末尾一次性解码**：生成完成后，把「上下文 + 生成 token」最后 512 个送进 tokenizer decoder，取最后 `pred_len` 行。 ### 4.3 多次采样取平均（抗噪核心） ```python x = x.unsqueeze(1).repeat(1, sample_count, 1, 1) # 复制 sample_count 条轨迹 ... preds = np.mean(preds, axis=1) # 多条轨迹平均 ``` 金融数据随机性大，单条采样路径抖动明显。`sample_count=N` 相当于蒙特卡洛估计： $$\hat{y} = \frac{1}{N}\sum_{i=1}^{N} f_\theta(\text{sample}_i)$$ 显然这是**用算力换稳定性**：耗时与 `sample_count` 近似线性增长。 ### 4.4 计算量 自回归部分每步要做一次全长前向（上下文 ≤ 512），因此总复杂度约为： $$O(\text{pred\_len} \times \text{sample\_count} \times L \times d_{model}^2 \times \text{ctx})$$ 这也是为什么仓库强烈建议 `lookback` 不要超过 `max_context`——超出的部分会被直接截断，白算。 --- ## 5. 与通用时序基础模型的差异 | 维度 | 通用 TSFM（如 TimesFM/Chronos） | Kronos | |---|---|---| | 输出空间 | 连续值回归 或 单层 VQ | **BSQ 二进制球面量化**，2^20 码字 | | 代码结构 | 单层 token | **s1(粗) + s2(细) 层级 token** | | 细粒度预测 | 独立头 | **依赖感知 cross-attention**（条件于 s1） | | 时间特征 | 往往只用 timestamp 编码 | minute/hour/weekday/day/month 相加 | | 抗噪手段 | 点估计 | **多次采样取平均**（sample_count） | | 领域先验 | 通用 | 6 维 K 线 + 交易时段 | 一句话概括：**Kronos 把「不确定性」显式建模进了架构**——粗码表达方向（高置信），细码表达幅度（低置信），采样平均抹平随机抖动。 --- ## 6. 代码地图 | 功能 | 位置 | |---|---| | BSQ 量化与熵损失 | `model/module.py::BinarySphericalQuantizer` (L39-223) | | 层级量化封装 | `model/module.py::BSQuantizer` (L225-255) | | RoPE / 注意力 / SwiGLU | `model/module.py::RotaryPositionalEmbedding` / `MultiHeadAttentionWithRoPE` / `FeedForward` | | 层级嵌入 | `model/module.py::HierarchicalEmbedding` (L400-444) | | 时间嵌入 | `model/module.py::TemporalEmbedding` (L536-563) | | 依赖感知层 | `model/module.py::DependencyAwareLayer` (L446-463) | | 双头 | `model/module.py::DualHead` (L486-514) | | Tokenizer | `model/kronos.py::KronosTokenizer` (L13-178) | | 语言模型 | `model/kronos.py::Kronos` (L180-328) | | 采样与自回归 | `model/kronos.py::top_k_top_p_filtering` / `sample_from_logits` / `auto_regressive_inference` | | 高层 API | `model/kronos.py::KronosPredictor` (L482-) | --- ## 7. 实测数据（本机 CPU，16 核） | 模型 | 参数量 | 加载(缓存) | 每步生成 | pred_len=70 总耗时 | |---|---|---|---|---| | Kronos-small | 24.7M | ~2.0s | 0.07–0.09s | ~5s（推算）| | Kronos-base | 102.3M | ~2.2s | 0.20–0.23s | 14.5s | （上下文 512，`T=1.0, top_p=0.9, sample_count=1`；`sample_count=5` 时约 ×5） --- ## 8. 局限与工程注意 1. **上下文硬上限**：small/base 只有 512，超过会被滑动窗口丢弃，长周期建模需要更大模型（mini 是 2048 但容量小）。 2. **采样成本**：`sample_count` 和 `pred_len` 都是线性开销，做批量回测时要先算好预算。 3. **tokenizer 与 LM 强耦合**：换数据集微调时，tokenizer 和 predictor 必须**成对**使用（仓库的 `finetune_csv` 流程也是先训 tokenizer 再训 predictor）。 4. **归一化在窗口内**：极端行情（涨跌停连续板）会让 `std` 失控，需要 `clip=5` 兜底。 5. **没有跨标的注意力**：一次只处理单条序列，组合层面的相关性不由模型建模。 --- ## 9. 参考 - 论文：*Kronos: A Foundation Model for the Language of Financial Markets*（arXiv:2508.02739） - BSQ：*Binary Spherical Quantization for Lightweight Image/Video Tokenization*（arXiv:2406.07548） - 本仓库：`model/module.py`、`model/kronos.py` - 微调细节：[`docs/finetune_guide.md`](finetune_guide.md)
avatar
文章
阅读
粉丝