SWARM-Lab/CoCo-Chronos2-LoRA

CoCo: Chronos2 LoRA adapter

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. A unified public CoCo inference package is being prepared.

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.