VITA: Variational Pretraining of Transformers for Climate-Robust Crop Yield Forecasting

Official implementation of VITA, a variational pretraining framework that learns weather representations from rich satellite data and transfers them to yield prediction tasks with limited ground-based measurements.

[arXiv:2508.03589] [Pretrained Model] [AAAI-26]

Overview

VITA addresses the data asymmetry problem in agricultural AI: pretraining uses 31 meteorological variables from NASA POWER satellite data, while deployment relies on only 6 basic weather features. Through variational pretraining with a seasonality-aware sinusoidal prior, VITA achieves state-of-the-art performance in predicting corn and soybean yields across 763 U.S. Corn Belt counties, particularly during extreme weather years.

Updates

• 11/21/2025: 🏆 VITA was selected for an oral presentation at AAAI 2026

Usage

Data Download

🛰️ Pretraining dataset: NASA POWER Daily Weather

🌽 Crop yield dataset: USA Corn Belt Crop Yield

pip install -r requirements.txt

python -m src.downloaders.nasa_power_dataset --data-dir path/to/data/dir
python -m src.downloaders.khaki_corn_belt_dataset --data-dir path/to/data/dir

Pretraining

python -m src.pretraining.main --batch-size 256 --n-epochs 100 --model-size small --alpha 0.5 --data-dir data/

Pretrained model weights: https://huggingface.co/notadib/VITA

Crop Yield Prediction

Note: This is an example run. For full hyperparameter configurations that reproduce paper results, see the paper's appendix. Due to non-determinism from hardware differences (GPU type, cuDNN versions) and stochastic training, you may observe small numerical variations from the exact values reported in the paper, though performance should remain in the same ballpark.

python -m src.crop_yield.main --batch-size 16 --n-epochs 40 --model-size small --beta 1e-4 --init-lr 2.5e-4 --test-type extreme --crop-type soybean --pretrained-model-path path/to/pretrained_model.pth

Citation

@inproceedings{hasan2026vita,
      title={VITA: Variational Pretraining of Transformers for Climate-Robust Crop Yield Forecasting},
      author={Adib Hasan and Mardavij Roozbehani and Munther Dahleh},
      booktitle={Proceedings of the 40th AAAI Conference on Artificial Intelligence},
      year={2026},
      url={https://arxiv.org/abs/2508.03589},
}