STAF-Net: An Innovative Framework for Wheat Yield Prediction

Abstract. Accurate crop yield forecasting is critical for optimizing agricultural resource management and ensuring food security. This study introduces STAFNet (Spatial-Temporal Attention Fusion Network), an innovative deep learning framework designed to integrate multispectral Sentinel-2 imagery and climatic variables for wheat yield prediction under limited data conditions. Classical machine learning models (Random Forest, XGBoost, Support Vector Machine) and a CNN-LSTM architecture were evaluated for comparison. Additionally, a Generative Adversarial Network (GAN) was employed to generate realistic synthetic multispectral images, addressing dataset scarcity and enhancing model generalization. Experiments were conducted in Sidi Yahya Zaer, Morocco, using simulated yield data derived from NDVI-based statistical modeling for the 2020–2024 period. Results show that XGBoost achieved strong baseline performance (R² = 0.919), while STAFNet exhibited superior temporal stability and accuracy. Incorporating GAN-based augmentation further improved STAFNet’s performance, reaching R² = 0.935 and significantly reducing RMSE and MAE. Multi-horizon testing confirmed robust early-season predictive capability from January onwards. These findings highlight the combined benefits of attention-based architectures and synthetic data generation for in-season yield forecasting, offering a scalable, cost-effective solution adaptable to various crops and regions.