CEiT — Contextually Extended iTransformer

Query-based virtual sensing for heterogeneous spatial sensor networks.

CEiT predicts the time series of any environmental variable at any location from an arbitrary set of physical sensor readings. Give it whatever sensors you have — temperature, humidity, CO₂, NH₃, at known 3D positions — and query a target (position, type); CEiT returns the predicted series there, even for a location and modality that was never instrumented. The sensor set can change at inference time with no retraining.

It extends the inverted Transformer (iTransformer) with three ingredients: a multi-modal token embedding (value + position + type), a hierarchical attention mask that lets query tokens attend to sensors but not vice versa, and self-reference training (SRT) for robustness to variable sensor configurations.

📄 Paper: Contextually Extended iTransformer (CEiT) for query-based virtual sensing in dynamic spatial sensor networks (see Citation below).

---

✨ Highlights

• Query-based. Predict at arbitrary (x, y, z) coordinates and sensor types, not just at fixed sensor locations.
• Set-flexible. Accepts any number of input sensors in any configuration; add/remove sensors without retraining or graph reconstruction.
• Heterogeneous fusion. Fuses multiple physical modalities (temperature, humidity, CO₂, NH₃) in a single model and supports cross-type prediction.
• Interpretable. The attention weights expose which sensors drive each prediction (used for sensor-importance and placement analysis).
• Strong accuracy. Up to ~50% lower RMSE than inverse-distance weighting (IDW) on real building data — see Results below.

What you can do out of the box: same-type virtual sensing, cross-type prediction, sensor-importance via attention, optimal sensor placement search, and continuous spatial-field reconstruction.

---

📦 Installation

git clone https://github.com/bet-lab/ceit.git
cd ceit
make setup          # == uv sync   (uses the uv package manager)

Requires Python ≥ 3.11 and uv. PyTorch ≥ 2.6 is installed from the CUDA 11.8 (cu118) wheel index; a GPU is recommended for training.

---

🚀 Usage

Run the model

import torch
from ceit import CEiT

# Instantiate with the paper configuration
model = CEiT(
    window=48,            # input length L (24 h at 30-min resolution)
    d_model=768,
    num_layers=3,
    nhead=8,
    value_mlp_layers=3,
    output_mlp_layers=2,
    dropout=0.27,
    use_gff=False,        # learnable linear position encoding (paper Eq. 3)
    use_type_embed=True,
)
model.eval()

# A batch with N physical sensors (any number) and Q queries (any location/type)
B, N, Q, L = 1, 9, 4, 48
batch = {
    "input_values":    torch.randn(B, N, L),          # sensor histories
    "input_positions": torch.rand(B, N, 3),           # normalized (x, y, z)
    "input_types":     torch.randint(0, 4, (B, N)),   # sensor type per input
    "input_mask":      torch.ones(B, N, dtype=torch.long),
    "query_positions": torch.rand(B, Q, 3),           # where to predict
    "query_types":     torch.randint(0, 4, (B, Q)),   # what to predict
    "query_mask":      torch.ones(B, Q, dtype=torch.long),
}

with torch.no_grad():
    out = model(batch)

print(out["query_values"].shape)   # -> (B, Q, L): predicted series at each query

Load a trained checkpoint and predict

from ceit import CEiT, CEiTPredictor

model, meta = CEiT.load_model("outputs/checkpoints/full_ceit/best_swa_model.pth", device="cpu")

# CEiTPredictor wraps batching + inverse normalization; build it with the
# dataset's normalizer and data feeder (see scripts/evaluate.py for the full setup):
#   predictor = CEiTPredictor(model, normalizer, data_feeder, device=...)
#   result = predictor.predict_single_window(
#       input_sensors=["T1", "T2", "H3", "C5"],   # IDs whose history feeds the model
#       query_positions=[(8500, 2225, 1900)],     # (x, y, z) in mm, or a known sensor ID
#       query_types=["TEMP"],                      # SensorType, name, or int
#       start_idx=0,
#   )
#   series = result["predictions"]                  # denormalized time series

Trained checkpoints are published as GitHub Release assets (tag weights-v1). Download full_ceit_best_swa_model.pth and place it at outputs/checkpoints/full_ceit/best_swa_model.pth; the ablation (no_type, no_srt) and Perceiver-IO variants are attached to the same release. Or reproduce from scratch with make train-all.

---

🏋️ Training

make train            # full CEiT          (configs/train.yaml)
make train-no-type    # ablation: no type embedding
make train-no-srt     # ablation: no self-reference training
make train-perceiver  # CEiT-P: Perceiver IO backbone swap
make train-all        # all four, sequentially (~3 h on a GPU)

# Optuna TPE hyperparameter search
make tune-ceit        # or: tune-perceiver / tune-all

Or call the script directly with config + CLI overrides:

uv run python scripts/train.py --config configs/train.yaml
uv run python scripts/train.py --config configs/train.yaml --model.d_model 512 --training.seed 123

Training uses SWA + early stopping and writes to a deterministic run_name directory under outputs/checkpoints/.

---

📊 Results

Same-type leave-one-out virtual sensing on real swine-farm data (held-out spring test period), CEiT vs. IDW interpolation:

Variable	CEiT RMSE	RMSE reduction vs IDW
Temperature	0.35 °C	40.6%
Humidity	1.33 %RH	20.5%
CO₂	145 ppm	50.6%
NH₃	0.12 ppm	— (tied; at sensor accuracy floor)

Full numbers — including the CEiT-P backbone-swap baseline and per-sensor breakdowns with standard deviations — are in the paper (Table 7). NH₃ readings sit near the sensor accuracy floor and should be read as a tied outcome rather than an improvement.

Evaluation & reproduction

make eval-all         # 5 tasks for full CEiT: same / cross / removal / placement / field
make eval-paper-all   # all tasks for all model variants
make all              # full pipeline: train -> eval -> stats -> curate

Curated metrics that back the paper live in results/ (git-tracked); raw run artifacts go to outputs/ (gitignored, regenerable). make verify runs the reproducibility test suite (checkpoint loading, RMSE-vs-paper consistency, paired statistics).

---

🗂️ Repository structure

├── src/ceit/                   # Core package
│   ├── model.py                # CEiT model
│   ├── data.py                 # SensorType, SensorDataset, DataFeeder
│   ├── predict.py              # CEiTPredictor, IDWPredictor
│   ├── perceiver.py            # CEiT-P (Perceiver IO backbone)
│   ├── metrics.py              # RMSE / MAE / comparisons
│   ├── optimization.py         # sensor-placement search
│   ├── interpolation.py        # spatial field reconstruction grid
│   ├── sensor_index.py         # physical-ID <-> spatial-label mapping
│   └── config.py               # dataclass configuration
├── scripts/                    # train.py, evaluate.py, tune.py, paired_stats.py, ...
├── configs/                    # YAML configs (train*/eval*/tune*)
├── dataset/SmartFarmData/      # sensor CSV data (git-tracked)
├── results/                    # curated metrics (git-tracked)
├── outputs/                    # checkpoints, eval dumps (gitignored)
└── tests/                      # checkpoint / consistency / statistics tests

---

⚙️ Configuration

Every experiment is fully specified by a YAML file plus optional --section.key value CLI overrides (see src/ceit/config.py). Each config pins a run_name so re-runs are idempotent. Key hyperparameters of the full model:

Symbol	Config key	Value
d_model	model.d_model	768
N_layer	model.num_layers	3
N_head	model.num_heads	8
dropout	model.dropout	0.27
L (window)	data.window	48
learning rate	training.learning_rate	1.3 × 10⁻⁴
p_self (SRT)	training.srt_prob	0.435

Position encoding is a learnable linear projection (use_gff: false); the Gaussian-Fourier path exists but is disabled. Determinism is on by default — for bit-exact retraining, export CUBLAS_WORKSPACE_CONFIG=:4096:8 before training.

---

🗃️ Data

dataset/SmartFarmData/ ships five months (Jan–Jun 2023) of indoor measurements from a commercial swine farm in Jeju, Republic of Korea — temperature, humidity, CO₂, NH₃, plus outdoor weather — as raw (*_1m-6m.csv) and 30-min resampled (*_resampled.csv) series, with sensor_mapping.csv and sensor_positions.csv.

---

📝 Citation

@article{choi2026ceit,
  title  = {Contextually Extended iTransformer (CEiT) for query-based virtual sensing in dynamic spatial sensor networks},
  author = {Choi, Wonjun and Lee, Sangwon and Shin, Hakjong and Kwak, Younghoon},
  year   = {2026},
  note   = {Under review; Code: https://github.com/bet-lab/ceit}
}

📄 License

Released under the MIT License. Maintainer: Sangwon Lee (lsw91.main@gmail.com).