Title: LoRA-EnVar: Parameter-Efficient Hybrid Ensemble Variational Assimilation for Weather Forecasting
Venue: NeurIPS 2025
Accepted: yes
Keywords: AI for science

data assimilation

weather forecasting

low rank adaptation
Abstract: Accurate estimation of background error (i.e., forecast error) distribution is critical for effective data assimilation (DA) in numerical weather prediction (NWP). In state-of-the-art operational DA systems, it is common to account for the temporal evolution of background errors by employing hybrid methods, which blend a static climatological covariance with a flow-dependent ensemble-derived component. While effective to some extent, these methods typically assume Gaussian-distributed errors and rely heavily on hand-crafted covariance structures and domain expertise, limiting their ability to capture the complex, non-Gaussian nature of atmospheric dynamics. In this work, we propose LoRA-EnVar, a novel hybrid ensemble variational DA algorithm that integrates low-rank adaptation (LoRA) into a deep generative modeling framework. We first learn a climatological background error distribution using a variational autoencoder (VAE) trained on historical data. To incorporate flow-dependent uncertainty, we introduce LoRA modules that efficiently adapt the learned distribution in response to flow-dependent ensemble perturbations. Our approach supports online finetuning, enabling dynamic updates of the background error distribution without catastrophic forgetting. We validate LoRA-EnVar in high-resolution assimilation settings using the FengWu forecast model and simulated observations from ERA5 reanalysis. Experimental results show that LoRA-EnVar significantly improves assimilation accuracy over models assuming static background error distribution and achieves comparable or better performance than full finetuning while reducing the number of trainable parameters by three orders of magnitude. This demonstrates the potential of parameter-efficient adaptation for scalable, non-Gaussian DA in operational meteorology.


Peer reviews:
comments: SUMMARY:
The authors develop a new variational assimilation method to support numerical weather prediction. First, a VAE is trained on historical forecast errors to serve as a baseline. Then, LoRA modules are inserted into the decoder. These LoRA modules are tuned at each time point using uncertainty information from AI-based ensemble forecasts. The resulting framework is intended to efficiently incorporate new information while keeping the historical background of the model.

STRENGTHS AND WEAKNESSES:
Strengths: The online training of LoRA modules in the decoder serves a double purpose to vastly improve efficiency over full fine tuning while preserving the decoder’s original parameters. Empirical results show strong improvement over previous approaches in a vitally important application.

Weaknesses: While the included spatial and temporal summaries are of greater interest and importance, it may also be interesting to see overall accuracy metrics to provide full context (perhaps just for your main results). While the incorporation of LoRA modules is novel, the overall VAE approach to data assimilation is familiar.

QUESTIONS:
The space/time summaries of forecast errors seem to favor your method, does this still hold up when reduced to a single spatial/temporal mean? Appears to be the case but it’s not always clear to me (especially in parts of Figure 3).

LIMITATIONS:
Yes.
TITLE: Official Review
CLARITY: 4
IMPACT: 4
ORIGINALITY: 3
DATE: 2025-06-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work addresses the problem of data assimilation by proposing a data-driven method based on the VAE-VAR architecture. A key contribution is the explicit separation of two types of model error: climatological error and flow-dependent error. These are handled differently: the former through offline training, and the latter via online fine-tuning of a small subset of network weights using the LoRA (Low-Rank Adaptation) technique. Experimental results show the effectiveness of both online adaptation in general, and the LoRA-based method in particular, compared to other fine-tuning approaches.\n\nIf I understood correctly, the core components of the proposed method are VAE-VAR and LoRA. The main novelty compared to existing machine learning-based approaches lies in the ability to adapt the prior model online through LoRA.\n\nSTRENGTHS AND WEAKNESSES:\nOverall, the paper is clearly written and the idea is original. However, I found some clarity issues and a general lack of context.\n\n### Strengths\n\n- Data assimilation is a crucial challenge, and leveraging deep learning has the potential to significantly improve weather and climate forecasting.\n- The paper is clearly written and easy to follow.\n- The proposed technique allows low-cost parameter adaptation.\n- The idea of adapting neural network parameters online is novel in this domain and appears both intuitive and powerful.\n- The experimental results are promising and support the main claims.\n\n### Weaknesses\n\n1. The paper focuses exclusively on spatial data assimilation, where observations are assimilated once per time step. There is no mention of 4D-Var, which is the state-of-the-art approach and would be a natural point of comparison.\n\n2. The distinction between climatological and flow-dependent error is not clearly defined. From my understanding, both arise from model forecast errors but differ in their time scales. Since this conceptual separation underpins the method, it would be valuable to define both types of error mathematically, for example, as expectations over model trajectories, and to clarify how they differ in structure and treatment.\n\n3. The paper lacks technical depth in describing its main components, VAE-VAR and LoRA. While the plain English explanations are helpful, the mathematical formulation is minimal. VAE-VAR is central to the model, yet it is only briefly described, with most details deferred to Appendix A. Similarly, LoRA is treated as well-known, which may not be the case in the scientific machine learning community. The problem of efficiently adapting a parametric model that defines a distribution, given samples from that distribution, is nontrivial and deserves formal treatment. This could benefit from connections to frameworks like online learning, adaptive neural networks, and meta-learning. While much of the existing literature in these areas focuses on prediction and forecasting tasks, rather than on estimating prior distributions, there are strong conceptual overlaps regarding online model adaptation. In particular, recent works on low-rank meta-learning architectures for physical systems [1, 2] may offer useful insights and parallels.\n\n4. Relatedly, the comparison to existing machine learning-based data assimilation methods (e.g., DiffDA) is somewhat superficial. A clearer description of the proposed architecture and its treatment of model error would enable more precise comparisons to alternative approaches. See also my Question 1.\n\n5. In the paragraph beginning at line 103, no references are provided for traditional hybrid EnVar approaches. A citation would strengthen the context and situate the contribution within established literature.\n\n### References\n\n[1] Serrano, L., Koupaï, A. K., Wang, T. X., ERBACHER, P., & Gallinari, P. Zebra: In-Context and Generative Pretraining for Solving Parametric PDEs.\n\n[2] Blanke, M., & Lelarge, M. Interpretable Meta-Learning of Physical Systems. In The Twelfth International Conference on Learning Representations.\n\nQUESTIONS:\n1. How does the proposed model compare with other architectures such as DiffDA? In line 45, the authors state: “but they do not incorporate climatological priors or flow-dependent ensembles.” Could the authors elaborate on this distinction and explain this mathematically, and how it translate into architectural or algorithmic differences?\n\n2. The training procedure uses the NMC method (mentioned in line 59 and described in the Appendix). Could the authors clarify how training error samples are constructed using this method, and how these samples capture 'climatological' rather than 'flow-dependent' error?\n\n3. Could the authors discuss the computational cost of their method? For instance:\n   - Does the VAE structure reduce cost by enabling assimilation in latent space?\n   - How does the LoRA-based online tuning compare in efficiency to other fine-tuning strategies?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper aims to address limitations in current variational data assimilation (DA) systems, most notably, their reliance on using a static background error learned from climatological data, and the limitation to using Gaussian errors to generate analysis. While hybrid DA methods exist that combine background error statistics with \"error-of-the-day\" statistics obtained by ensemble simulations, traditionally, this still assumes the error to be Gaussian to be used within the traditional variational framework. The paper proposes an algorithm that performs variational data assimilation in latent space using VAEs trained on climatological error samples; assuming a Gaussian bottleneck for the VAE, the use of Gaussian error statistics is justified in latent space by construction. Furthermore, the VAE is fine-tuned using LoRA on \"errors-of-the-day\" to achieve a dynamically evolving error statistics. Experiments on a weather forecasting model trained on ERA5 demonstrate the added benefits of using non-Gaussian error estimates (achieved by operating in latent space) and dynamically evolving error by LoRA finetuning.\n\nSTRENGTHS AND WEAKNESSES:\n__Strengths:__\n- Overall, I believe that the idea presented in this work is very elegant, providing an interesting approach to tackle both the issues of using static and Gaussian error covariances in traditional data assimilation using a single clean framework.\n- The experiment section is quite strong. In particular, the ablation studies convincingly demonstrate the added benefit of using both non-Gaussian error statistics and dynamical error statistics in data assimilation, at different observation densities.\n- Experiments to compare LoRA vs full fine-tuning demonstrate that we are not losing much by only fine-tuning the VAE; in fact, the experiments show that LoRA is not only sufficient, but even outperforms the results of full fine-tuning. in general, I believe that all of the components that are introduced in the model is justified.\n\n__Weaknesses:__\n- One clear limitation of the approach is that the method explicitly uses the fact that the VAE has a Gaussian bottleneck to justify the use of 3DVAR in latent space. Thus, it cannot accommodate the use of more general autoencoders such as vector-quantised VAEs (VQ-VAEs), which can achieve better reconstruction compared to standard VAEs.\n- A minor gripe that I have is that some of the colors used in the plots are very similar to each other, which makes it hard to see which curve corresponds to which model.\n- I believe the proposed methodology can be better explained. For example, I believe it would be useful to have Appendix A in the main body instead of the appendix. Otherwise, we are only left with a vague description of the methodology and Figure 1 to understand the idea, which I initially found difficult to grasp until reading Appendix A.\n\nQUESTIONS:\n- How are the ensemble obtained to fine-tune LoRA (i.e., what are the \"Other predictions\" in Figure 1)? Are they based on different initial conditions or different forecasting models?\n\nLIMITATIONS:\nThe authors acknowledge limitations of the approach, in particular, lack of theoretical characterisation and intrinsic limitation of only being able to produce point estimates.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe submission introduces a new data assimilation approach which uses a pretrained VAE to learn an approximation of a non-Gaussian background covariance which is then updated online via LoRA in order to incorporate information from the local flow. This approach is empirically validated against several traditional approaches and the recently proposed VAE-Var on both assimilation from reanalysis ICs and simulated operational ICs across a variety of observation densities.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n- S1 - The challenge of efficiently capturing non-Gaussian variation in data assimilation is a key challenge in modern weather forecasting. \n- S2 - The combination of experiments progressively challenges the model. The use of restricted observations in space and time are very important and relevant to practical usage and greatly strengthen the potential impact. \n- S3 - The extensive ablations are really interesting and provide some somewhat counter-intuitive information (LoRA providing the biggest gains for slower moving fields like Z500 in Fig 3, larger amounts of finetunable parameters hurting performance). This could be explored more strengthening the paper. \n- S4 - Performance at smaller ensemble sizes is extremely promising. In an operation setting this could significantly reduce the cost of the forecast.\n\nWeaknesses:\n- W1 - The major weaknesses of the paper are centered on the evaluation and comparisons.\n    - W1.A - The comparisons in general are very weak. The proposed method utilizes a 300M parameter SWINv2 architecture while 3DVar is largely a teaching example. This is not a fair comparison. While it's not expected that the experiments use a full operational DA system, a variational-to-variational comparison should include 4DVar (and associated ensemble-equipped methods) and likely at least one EnKF-based method (as they tend to be a bit easier to tune). Given the small ensemble sizes, it's likely that these would require some parameter tuning.\n    - W1.B - On the ML side, the related work discussed a number of different approaches while stating without evidence why these approaches are unsuitable. This is a largely empirical submission, so these claims need to be shown rather than just stated. If scale is an issue, it's fine to just point to an appendix section showing why those approaches shouldn't scale. If the issue is instead that methods perform poorly at scale, then this can be shown experimentally. \n    - W1.C - Right now the paper is motivated by fixing an error in the probabilistic formulation of variational DA. However, there is then no follow-up on calibration of forecasts. The paper states that this is not available for variational methods which is true for pure variational approaches where all spread information is provided a priori, but with hybrid methods, this can be estimated empirically from the ensemble.\n - W2 - All experiments are currently performed on a fairly small time window. 4.2 is on January 1, 2022. 4.3 is on that same month. There does not seem to be a reason why more extensive backtesting would be impossible. \n- W3 - (Minor) Parameter efficiency is interesting, but not really the relevant comparison given your experiments show LoRA outperforming full rank training. Since online training is involved, LoRA should actually provide some real hardware efficiency gains either on runtime for fixed hardware or memory usage. These are more relevant.\n- W4 - (Minor) For cases where a term or definition is vital to understanding the paper, it's usually better to find a way to include it in the main text rather than pointing readers to citations or the appendix. The variational procedure is the most notable case of this in the text currently.\n\nOverall, this feels quite close to the line. Maintaining results with nontrivial numerics baselines and performing more extensive testing across time is likely enough that I'd be comfortable moving to weak accept while including stronger ML baselines, including analysis of the field distributions (appendix is fine as long as it's pointed out and discussed in the main body), and more probabilistic analysis would result in a stronger evaluation.\n\nQUESTIONS:\nQ1 - How does the calibration of the ensembles resulting from this approach compare to those of the competitive approaches?\n\nQ2 - It is often easy to reduce RMSE by reducing high-frequency details and extremes in the forecast. Does the distribution of the resulting field values match the target?\n\nQ3 - Does this significant competitive advantage persist when compared against modern DA tools?\n\nLIMITATIONS:\nNegative societal impact yes. However, discussion of limitations largely focuses on trivial limitations without discussing limitation in experimental methodology.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors use finetuned LoRA modules and combine them with climate error models (as in VAE-VAR) to conduct data assimilation. They do this by first training a VAE decoder on historical error values and then creating an ensemble of predictions for each cycle and finetuning the LoRA matrices for the flow-dependent ensemble perturbations.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\nThe model is efficient in terms of trainable parameters at test time which is only 0.1 % of the model, which is practical for finetuning with limited GPU capacity when running in an online fashion. This allows one to adapt it to more complex transformer architectures for the VAE without much computational overhead.\n\nThe authors do not assume a Gaussian background which allows for more flexible likelihood functions compared to standard data assimilation models which assume Gaussian backgrounds.\n\nThe authors test the model at 0.25 degree global resolution, and they found that LoRA with rank 2 was sufficient to capture the detailed dataset, so method scales without exponential parameter growth. 0.25 degree forecasting is the de-facto standard for MLWP today.\n\nThe model also seems robust to different ensemble sizes, observation densities, and variables. The authors have done quite a few ablation experiments while adjusting these parameters.\n\nThe integration of LoRA into standard diffusion model transformers for image generation has been explored before, but it has not been applied to the weather prediction and assimilation community. This paper presents the first foray into incorporating such a step, and also offers avenues of improvements for future data assimilation techniques which are based on a VAE-Var style of data assimilation. \n\nLORA EnVar seems to scale well with observation sparsity in comparison to the baseline methods. However, the results are not a significant boost over strong baselines such as VAE-Var.\n\nWeaknesses\n\nThe model relies on a Swin-V2 backbone tailored to FengWu; portability to other forecast models is not clear in terms of how it goes.\nAlthough LoRA updates are light, the decoder is still a large transformer that must run in each inner-loop gradient evaluation.\n\nSummary\n\nThe paper is well-structured and clear, and offers significant novelty in terms of introducing LoRA to data assimilation with MLWP, and opens up future avenues of improvement for VAE-Var style data assimilation approaches. The results seem like a modest improvement upon existing data assimilation techniques like VAE-Var. As a result, I vote accept for this paper.\n\nQUESTIONS:\nWhy is are only the query LoRA integrated, as compared to also the K and the V?\n\nIs there a link between the LoRA rank and an assumed effective ensemble size or background covariance rank?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}