Lucifer744/KiT_model

KiT Model Checkpoints

Pre-trained model checkpoints for KiT: A Foundation Model for Candlestick Time-Series Forecasting via Diffusion Transformers.

Checkpoint bundles

Each checkpoint is stored with its matching instrument mapping and normalization statistics:

step_market_284M/
  step_market_284M.pt
  instruments.json
  stats.json
step_market_284M_train260410/
  step_market_284M_train260410.pt
  instruments.json
  stats.json

Bundle	Markets	Instruments	Normalization entries	Training cutoff
step_market_284M	Chinese A-shares, U.S. equities, cryptocurrency	4,830	105	2025-12-31
step_market_284M_train260410	Chinese A-shares	3,914	35	2026-03-09

step_market_284M

The official open-source checkpoint used for the paper's reported results and evaluations. Its EMA weights exactly match the original three-market experiment's step 64,000 checkpoint. The existing .pt file retains a legacy step=50000 metadata field; the weight tensors are unchanged.

step_market_284M_train260410

The A-share checkpoint whose EMA weights exactly match the original experiment's step 63,000 checkpoint. The filename is retained for compatibility. The original training configuration and launch records specify training through 2026-03-09, with 2026-03-10 through 2026-04-10 held out for validation; the filename does not mean that training included all data through April 10. This checkpoint has no reported test-set benchmark.

Using the matching metadata

• instruments.json maps each symbol to its market, integer sector, and original vocab_index. Preserve these IDs: they select the trained instrument embeddings.
• stats.json contains the normalization scales used for input normalization and output denormalization, keyed by <market>|<timescale>|<feature>.
• For the three-market bundle, A-share and U.S. statistics come from their original stats.json files; cryptocurrency statistics come from the original stats_trainend2025.json, as selected by the training merge code.
• These files were reconstructed from the available original market metadata using the training merge logic, without renumbering instruments or refitting statistics. Eight real-data inference cases across both checkpoints were compared against the matching original experiment weights and source metadata using identical seeds: normalized predictions and reconstructed OHLCV matched exactly. This verifies compatibility on the tested cases; training-time JSON hashes were not recorded.

Download and inference

Download either bundle, for example:

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="Lucifer744/KiT_model",
    allow_patterns=["step_market_284M/*"],
    local_dir="ckpt",
)

The checkpoint path is now ckpt/step_market_284M/step_market_284M.pt. For the A-share-only bundle, select step_market_284M_train260410/* and use the corresponding nested checkpoint path.

The current code repository reads instruments.json and stats.json from --data-root. Place the chosen bundle's two JSON files in that data directory alongside your market/timeframe parquet files and trading calendar, then pass the nested checkpoint path explicitly:

python eval/infer_for_pred.py --ckpt ckpt/step_market_284M/step_market_284M.pt --data-root data

The bundled evaluation scripts target A-shares. Use the lower-level inference API with the appropriate market and calendar for U.S. equities or cryptocurrency. Keep the metadata from the same bundle as the selected checkpoint.

The 100-stock A-share demo data is available at KiT_data_demo.

Citation

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

License

See the main repository for license information.