CoCo: Chronos2 LoRA adapter

Paper · Blog · Code

The multivariate time-series forecasting component of TSFM Meets LLM: Context as Covariates (CoCo), shared by SWARM-Lab. This repository contains the adapter from the paper's submitted supplementary package, not the base model weights.

This adapter teaches Chronos2 to use an LLM's forecast and confidence covariates together with numerical history. Its companion LLM adapter is SWARM-Lab/CoCo-Qwen3-4B-LoRA.

Checkpoint

Base model: amazon/chronos-2.
PEFT LoRA rank: 8; alpha: 16; dropout: 0.
Target modules: self_attention.q, self_attention.k, self_attention.v, self_attention.o, output_patch_embedding.output_layer.
Adapter file: 4,854,584 bytes.
Adapter configuration records PEFT 0.18.1.

Loading the component

import torch
from chronos.chronos2 import Chronos2Model, Chronos2Pipeline
from peft import PeftModel

base = Chronos2Model.from_pretrained(
    "amazon/chronos-2", torch_dtype=torch.bfloat16, device_map="cuda"
)
model = PeftModel.from_pretrained(base, "SWARM-Lab/CoCo-Chronos2-LoRA")
model = model.merge_and_unload()
model.eval()
pipeline = Chronos2Pipeline(model)

Merging happens in memory; this repository stores only the LoRA adapter. This follows the submitted loader, but the remote-loading example has not yet been GPU-tested as part of the release.

Use the paper's covariate construction rather than supplying arbitrary text embeddings: the forecast channel is shifted, numerical history is normalized using history-only statistics, and raw rollout variance is divided by squared history standard deviation (with the implementation's numerical stabilizer). The submitted implementation uses history noise standard deviation 0.05. The unified CoCo inference package provides the full pipeline with Transformers + PEFT or vLLM.

The submitted environment pins chronos-forecasting==2.1.0, peft==0.18.1, and torch==2.9.0.

Training and intended use

The paper describes second-stage LoRA adaptation on pre-generated LLM forecast and confidence covariates. Intended use is research on text-aided time-series forecasting. The confidence covariate is a rollout-variance signal, not a calibrated probability of correctness.

Base-model license

The base Chronos2 checkpoint is distributed under Apache-2.0 by its publisher.

Citation

If you use CoCo in your research, please cite:

@inproceedings{ram2026coco,
  title = {{TSFM} Meets {LLM}: Context as Covariates},
  author = {Hemanth Ram and Sai Shankar Narasimhan and Shubhankar Agarwal and Sandeep Chinchali},
  booktitle = {The Fortieth Annual Conference on Neural Information Processing Systems},
  year = {2026},
  url = {https://openreview.net/forum?id=t2nNTWgHfq},
  arxiv = {https://arxiv.org/abs/2611.05515}
}