KiT: A Foundation Model for Candlestick Time-Series Forecasting via Diffusion Transformers

KiT is a diffusion-based foundation model for candlestick (K-line) forecasting. It is trained on billions of bars spanning U.S. equities, Chinese A-shares, and cryptocurrencies across seven granularities, from one minute to one day, and achieves SOTA performance in both return forecasting and volatility prediction.

📑 Todo List

• Train
  • training infra
  • training code
• Eval
  • specific stock forecasting
  • inference over the validation set
  • simple backtest strategy
• Data
  • 100 stocks demo data
  • data preparation scripts
• Model
  • KiT-L: step_market_284M.pt
  • KiT-L: step_market_284M_train260410.pt

📖 Intro

KiT casts multi-horizon candlestick forecasting as conditional path generation via flow matching. The overall pipeline is illustrated above: raw OHLCV bars are encoded into a five-dimensional log-ratio state $x_t=(r_{\mathrm{gap}}, r_{\mathrm{body}}, r_{\mathrm{up}}, r_{\mathrm{dn}}, v_t)$, which is the state the diffusion model operates on. History and horizon are assembled into a single token sequence and processed by the KiT backbone, the history is returned bit-identical and only the forecast span is filled in with generated bars. In KiT block, signals that are constant over the window modulate every layer through a shared AdaLN trunk, whereas signals that vary per bar are added directly to the token embeddings.

💹 Prediciton Demo

Backtest

Return & volatility forecasting

Method	1m ↑	5m ↑	15m ↑	30m ↑	1h ↑	2h ↑	1d ↑	Mean ↑
KiT	.0215	.0576	.0526	.0474	.0383	.0651	.1168	.0571
Kronos-base	.0158	.0426	.0442	.0064	.0128	.0514	.1046	.0342
Kronos-base-FT	.0192	.0536	.0484	.0408	.0304	.0622	.1101	.0521
Sundial	.0148	.0386	.0434	.0082	.0046	.0214	.0837	.0209
Chronos-Bolt extended	.0362	.0088	.0186	.0964	.0068	.2159	.0052	.0461
Chronos-2	.0168	.0286	.0362	.0184	.0162	.0648	.0556	.0054
TimesFM 2.5	.0126	.0284	.0160	.0797	.0089	.0747	.0116	.0181
TimesFM 3	.0046	.0338	.0374	.0168	.0294	.0586	.0479	.0111
Historical return bootstrap	.0296	.0128	.0346	.0267	.0341	.0020	.0982	.0120
TimeGrad adapted	.0136	.0088	.0362	.0046	.0248	.0574	.0792	.0250
Diffusion-TS adapted	.0090	.0284	.0206	.0292	.0746	.0198	.1654	.0190
TSFlow adapted	.0174	.0554	.0466	.0228	.0102	.0596	.0673	.0399

Table 1. RankIC of Return forecasting.

Method	1m ↑	5m ↑	15m ↑	30m ↑	1h ↑	2h ↑	1d ↑	Mean ↑
KiT	.7599	.6947	.6900	.5919	.5958	.6652	.6225	.6600
Kronos-base	.5900	.4934	.3955	.3049	.4699	.5281	.4498	.4617
Kronos-base-FT	.6193	.5478	.4466	.3671	.4974	.5306	.4727	.4974
Sundial	.6024	.5477	.5164	.4591	.5691	.5534	.4486	.5281
Chronos-Bolt extended	.5766	.4821	.4113	.4289	.3543	.3797	.3310	.4234
Chronos-2	.4427	.3944	.3671	.4092	.4264	.4452	.2260	.3873
TimesFM 2.5	.5184	.4609	.4435	.4537	.5407	.4688	.3096	.4565
TimesFM 3	.3870	.4655	.4117	.3474	.4470	.3977	.2507	.3867
GARCH	.3520	.5940	.6469	.4850	.5842	.5681	.4255	.5222
TimeGrad adapted	.0736	.1967	.0056	.1586	.5201	.4418	.1702	.1785
Diffusion-TS adapted	.4462	.1787	.1483	.0246	.4236	.3488	.4563	.2895
TSFlow adapted	.3244	.3230	.4043	.5658	.5290	.3648	.5250	.4338

Table 2. RankIC of Volatility prediction

🚀 Get started

1. Setup

Requirements

• A CUDA GPU for inference (6GB+ VRAM recommended)
• Python >= 3.10
• PyTorch >= 2.4 (with CUDA support)

Installation

1. Clone the repository:

git clone https://github.com/Luciferbobo/KiT.git
cd KiT

2. Create and activate a conda environment:

conda create -n kit python=3.13 -y
conda activate kit

3. Install dependencies:

# Install PyTorch with CUDA support (adjust cuda version as needed)
pip install torch>=2.4 --index-url https://download.pytorch.org/whl/cu118

# Install other requirements
pip install -r requirements.txt

Download Model and Data

You can download the pre-trained model checkpoint and demo data from:

	Location	Hugging Face Link
Demo data	data/	KiT_data_demo
Model checkpoint	ckpt/	KiT_model

Currently, there are two versions of the KiT available. See KiT_model for more details.

Model	Training range	Ckpt Name
KiT-L	2016-01-01 to 2025-12-31	step_market_284M.pt
KiT-L	2016-01-01 to 2026-04-10	step_market_284M_train260410.pt

Quick download commands:

# Install huggingface-cli if not already installed
pip install huggingface_hub

# Download the model checkpoint
huggingface-cli download Lucifer744/KiT_model --local-dir ckpt/

# Download the demo data
huggingface-cli download Lucifer744/KiT_data_demo --local-dir data/ --repo-type dataset

Using Your Own Data

The demo data includes 100 stocks. If you want to use your own data instead:

1. Data format: Each stock needs parquet files per timescale (1m, 5m, 15m, 30m, 1h, 2h, 1d) with columns:

   • date: timestamp (datetime64[ns])
   • open, high, low, close: prices (float64)
   • volume: trading volume (float64)

2. Directory structure:

   data/
   ├── instruments.json       # List of instruments with metadata
   ├── stats.json            # Normalization statistics
   ├── trade_dates.csv       # Trading calendar
   └── ashare/               # Or your market name
       ├── 000001.SZ_1m.parquet
       ├── 000001.SZ_5m.parquet
       └── ...
   

3. instruments.json should contain:

   {
     "000001.SZ": {
       "name": "xxxx",
       "market": "SZ",
       "sector": "xxx",
       "list_date": "1991-04-03"
     }
   }

Refer to the demo data structure for the exact format requirements.

3. Train

TBD

4. Eval

(1) Forecast a specific stock and time

A case is code,timescale,anchor. The anchor is the end of the history window: the model reads the most recent Lc bars up to the anchor and predicts the next Lh bars.

python eval/infer_for_pred.py
python eval/infer_for_pred.py --case 601899.SH,5m,2026-03-04 --case 300750.SZ,30m,"2026-03-10 11:30"

Outputs go to results/infer_for_pred/. Plot legend: grey candles = tail of the history; red/green = real future; blue/purple = representative predicted path; blue bands = q05–q95 and q25–q75 of the K paths' close.

Window sizes are fixed by the model:

Timescale	1m	5m	15m	30m	1h	2h	1d
History Lc	1200	960	800	640	400	360	250
Horizon Lh	120	96	64	40	20	20	20

We provide visualization for the prediction results. Your prediction results will be automatically plotted, and it should look like this:

(2) Inference over the validation set

Runs all 100 stocks, all 7 timescales, all anchors in the val window (2026-01-01 to 2026-04-10)

python eval/infer_for_eval.py
python eval/cal_metrics.py

Task	Prediction	Target
return	mean over K paths of the sum of per-bar log returns	sum of real per-bar log returns
vol	mean over K paths of the std (ddof=1) of per-bar log returns	std of real per-bar log returns
price	mean over K paths of the cumulative log return vs. anchor close	real cumulative log return

(3) 💰 Backtest

After inference over the whole validation set has finished (results/infer_for_eval/raw/ is filled), if you are interested in real-world trading, we provide a backtest framework. Run the backtest directly on it:

python eval/backtest.py
python eval/backtest.py --top 20 --slip 5 --scales 5m,15m,1h

Rules:

• Signal: at every anchor, rank the stocks by the mean over the K sampled paths of the predicted window return.
• Portfolio: hold the top N with equal weight from the anchor close to the end of the forecast window, then rebalance at the next anchor. Windows do not overlap, so the returns of consecutive windows are chained. Stocks that stay in the top N are kept without trading.
• Limit-up / limit-down: a stock that closes limit-up at the anchor (or is in its first 5 listing days) cannot be bought and is skipped. A stock that is limit-down at the last bar of the window cannot be sold and is carried into the next window. Limits are 10% (20% for 300/301/688), a stock within 0.5% of the limit counts as sealed. ST stocks (5%) are not modelled.
• T+1: positions are bought at the close and sold in a later window, so it always holds.
• Costs: commission 2.5bp and transfer fee 0.1bp per side, stamp tax 5bp on sell, plus slippage (10bp per side by default). Costs are charged on the traded amount only. Lot size and market impact are ignored.
• Benchmarks over the same period: equal-weight mean and median of the 100 stocks per window (no costs), and CSI500 (daily closes, downloaded once with akshare and cached in data/csi500_daily.csv; if the download fails the benchmark is left out).

Results are written to results/backtest_results/

---

Citation

If you use KiT in your research, please cite:

@article{kit2026,
  title={KiT: A Foundation Model for Financial Time-Series Forecasting using Diffusion Transformers},
  author={Boyu Zhang, Haorui Li},
  journal={arXiv preprint arXiv:2609.34507},
  year={2026}
}

💬 Community

Welcome to join the KiT WeChat group! If the group QR code has expired, add the assistant and we'll invite you.

    
_{WeChat Group  |  Assistant}

License

Released under the AGPL-3.0 license.