Predictive Differential Training Guided by Training Dynamics
Authors: Fanqi Wang
Weisheng Tang
Landon Harris
Hairong Qi
Dan Wilson
Igor Mezic
Venue: ICLR 2026 Poster (2026)
Keywords: Training Dynamics
Koopman Operator Theory
Predictive Training
Deep Neural Networks
Primary area: optimization

Abstract: This paper centers around a novel concept proposed recently by researchers from the control community where the training process of a deep neural network can be considered a nonlinear dynamical system acting upon the high-dimensional weight space. Koopman operator theory (KOT), a data-driven dynamical system analysis framework, can then be deployed to discover the otherwise non-intuitive training dynamics. Taking advantage of the predictive power of KOT, the time-consuming Stochastic Gradient Descent (SGD) iterations can be then bypassed by directly predicting network weights a few epochs later. This "predictive training" framework, however, often suffers from gradient explosion especially for more extensive and complex models. In this paper, we incorporate the idea of "differential learning" into the predictive training framework and propose the so-called "predictive differential training" (PDT) for accelerated learning even for complex network structures. The key contribution is the design of an effective masking strategy based on a dynamic consistency analysis, which selects only those predicted weights whose local training dynamics align with the global dynamics. We refer to these predicted weights as high-fidelity predictions. DT also includes the design of an acceleration scheduler to adjust the prediction interval and rectify deviations from off-predictions. We demonstrate that PDT can be seamlessly integrated as a plug-in with a diverse array of existing optimizers (SGD, Adam, RMSprop, LAMB, etc.). The experimental results show consistent performance improvement across different network architectures and various datasets, in terms of faster convergence and reduced training time (10-40%) to achieve the baseline's best loss, while maintaining (if not improving) final model accuracy. As the idiom goes, a rising tide lifts all boats; in our context, a subset of high-fidelity predicted weights can accelerate the training of the entire network!

Published as a conference paper at ICLR 2026
PREDICTIVE DIFFERENTIAL TRAINING GUIDED BY TRAINING DYNAMICS
Fanqi Wang': Weisheng Tang'; Landon Harris', Hairong Qi', Dan Wilson', Igor Mezic² University of Tennessee, Knoxville UUniversity of California, Santa Barbara [fwang20, 1harri73]evoll edd, (wtang7 hgi, wwwdwwwloooo.oo.. lutk. edu, meziclucsb. edu
ABSTRACT
This paper centers around a novel concept proposed recently by researchers from the control community where the training process of deep neural network can be considered a nonlinear dynamical system acting upon the digensiimnsioall weight space. Koopman operator theory (KOT), data-driven dynamical system analysis framework, can then be deployed to discover the otherwise non-intuitive training dynamics. Taking advantage of the predictive power of KOT, the tine-conuuming Stochastic Gradient Descent (SGD) iterations can be then bypassed by directly predicting network weights few epochs later. This "predictive training" frame- work, however, often suffers from gradient explosion especially for more extensive and complex models. In this paper, we incorporate the idea of "differential learn- ing" into the predictive training framework and propose the so-called "predictive differential training" (PDT) for accelerated learning even for complex network structures. The key contribution is the design an effective masking strategy based on a dynamic consistency analysis, which selects only those predicted weights whose local training dynamics align with the global dynamics. We refer to these predicted weights as high-fidelity predictions. PDT also includes the design of an acceleration scheduler to adjust the prediction interval and rectify deviations from off off-predictions. We demonstrate that PDT can be seamlessly integrated as a plug-in with a diverse array existing optimizers (SGD, Adam, RMSprop, LAMB, etc.). The experimental results show consistent performance improvement across different network architectures and various datasets, in terms of faster convergence and reduced training time (10-40%) achieve the baseline's 's best loss, while main- taining (if not improving) final model accuracy. As the idiom goes, a rising tide lifts all boats; in our context, subset of high-fidelity predicted weights can accelerate the training of the entire network!
INTRODUCTION
The advent of cutting-edge hardware (Li et al. 2014) and the development of parallel processing tech- niques (Li et al., 2020) have greatly accelerated the training process of Deep Neural Network (DNN). However, enhancing the fundamental techniques of DNN training continues to be significant chal- lenge. From the inception of Stochastic Gradient Descent (SGD) (Robbins & Monro 195T), which has since become mainstay in DNN training, numerous techniques have been proposed to increase the efficiency of the underlying optimization task, including, for example, learning rate annealing and momentum (Sutskever al. 2013)), RMSprop (Tieleman & Hinton 2012), and Adam (Kingmal & Ba 2014). In addition to these first-order optimizers, second-order alternatives (Martens (0100) utilizing curvature information or second-order derivatives of the loss function have been explored to potentially enable more efficient convergence. Despite these ad advancements, gradient- based meth- ods are still inherently iterative, requiring repeated gradient computations and weight adjustments throughout the network. This iterative burden manifests fundamental limitation of SGD and its variants, which lies the core of the computationally expensive training process.
"Equal contribution

---

predictions Published as a conference paper at ICLR 2026
The concept of differential learning-wee wheee different parts of the network exhibit different learning behaviors during training has emerged as a promising direction to address this limitation. Differential learning can be layer-specific (Devlin et al., 2019; He et al.] 2019) or parameter-spectte (Tieleman & Hinton, 2012: Duchi et al. 201 Ia), allowing for more targeted optimization. The Adam optimizer (Kingma & Ba 2014), for instance, adaptively computes individual learning rates for different parameters, While differential learning takes adaptive approaches on how parameters are updated, it does not fundamentally address the limitation of the iterative optimization process itself.
Projection High Dimensional Loss Landscage
Space
Xzam Oura
Foo log
1.5-
Epect=10. Los02.03
0.5
0.5
Etoch-10
Losc-2 16
C0sch=20 Loaue
-15
RC1
Figure 1: Comparison of training trajectories and loss landscapes between Adam and the proposed PDT. (Trained on CIFAR-10 using Alex
Recently, a novel interpretation of the DNN training process has been proposed, mainly by researchers
from the control community (Redman et al., 2022, Dogra & Redman 2020 Manojlovic et al. 2020
Tano et al. 2020; Redman et al.] 2024) it intuitive to consider a pre-trained DNN as an inherently nonlinear static system acting upon the inputs, then the DNN "training process" itself is a "nonlinear"
dynamical system acting upon the high-dimensional "weight space"! It is a discrete dynamical system since the weights of a DNN evolve over each iteration (or epoch) according to the optimization process adopted. This drastically different interpretation has led to the establishment of novel mathematical framework for learning. Koopman Operator Theory (KOT) (Mezić 2005). a powerful data-driven dynamical system analysis tool, is often adopted to exploit the underlying dynamics in the seemingly non-intuiie ann-inuitiiee training process of a DNN. Taking advantage of the predictive power of KOT, the time-consuming SGD iterations can be bypassed by directly predicting network weights
a few epochs later (Dogra) 22020, Dogra & Redman 2020) Tano et al., 2020). We refer to these approaches as predictive training.
However, practical challenges quickly emerge. The absence of actual gradient descent means that convergence cannot be guaranteed, and the framework is sensitive to disturbances in the weight space, leading to error accumulation across iterations. As the network scales, the previous Koopman-baat bassed predictive training framework becomes increasingly ineffective. This issue is mostly due to the lack of adaptive mechanisms when applying prediction-baset acceleration. That is, existing predictive training approaches tend accept all predicted weights without checking if the prediction is of "high-fidelity" or not. This often leads to gradient explosion, especially for more extensive and complex models.
The key observation is that even though KOT is a powerful predictive tool for studying traditional small-scale control problems, when dealing with DNN whose parameter dimension reach into the millions or even billions, the quality of prediction tends to be highly inhomogeneous across the entire weight space. Hence, the predictive learning has to be "selective" only should be selected to effectively accelerate learning.
In this paper, we propose predictive differential training (PDT), where acceleration by prediction is selectively applied based on dynamic consistency analysis. This principled approach identifies parameters that are in a stable, predictable phase of their evolution by ensuring their local dynamics align with the global system dynamics modeled from the training history. This selective acceleration is conceptually similar to various adaptive learning rate methods. For instance, Adagrad (Duchi et al. 2011b) targets acceleration at rare features; Momentum (Rumelhart al.) 1986) prioritizes weights with the largest recent velocity; and the popular Adam optimizer (Kingma & Bal 2014) employs a combined strategy. Fig. illlustrates the compelling effectiveness of PDYT over Adam through visual
2

---

Published as a conference paper at ICLR 2026
comparison of the training trajectory and loss landscape. The contributions of the proposed PDT can be summarized as follows:
We propose the Predictive Differential Training (PDT) framework that selectively applies predictive updates to effectively accelerate training.
We design a dynamic consistency analysis as masking strategy to conduct prediction. selects parameters whose local training dynamics align with the global dynamics, allowing PDT to identify parameters that are stable, predictable phase of their evolution.
We demonstrate that PDT can be seamlessly integrated as a plug-in with numerous existing optimizers, such as SGD, Adam, RMSprop. Shampoo, and LAMB, while maintaining computational efficiency through epoch-level predictions.
We validate PDT's effectiveness across diverse network architectures (from FCN to ViT- Huge), datasets (from CIFAR-I0 10 to ImageNet), and learning paradigms (from supervised to self-suuprvised), demonstrating its scalability and robustness under various conditions.
2 BACKGROUND AND RELATED WORK
The key notion of Koopman analysis is the representation of a (possibly nonlinear) dynamical system linear operator on typically infiite-dimennioonll space of functions (Mezić 2021 2005 Mezić & Banaszuk 2004 Wang et al. 2024). Koopman-basee approaches directly contrast with standard linearization techniques that consider dynamics around nominal solution. Indeed, Koopman analysis can yield linear operators that accurately capture fundamentally nonlinear dynamics.
Koopman Operator Theory. As brief description, consider a discrete-time dynamical system Xi+1 = Txx)). where x; E R" is the current state and X is the next state after applying the potentially nonlinear mapping T. Consider also vector-valued observable g(x) E IR.". The evolution of observables under this mapping can be described as
g(x++1) = g(T())) Kg(x.).
(1)
where K. operates the vector space of observables and maps g(x;) to gxx+1)) K. is referred as the "Koopman operator" that is associated with the fully nonlinear dynamical system. The Koopman operator is linear, but also spectrum for observables MMezid 2020), its action can be decomposed according to
As such, for dynamical systems with pure point
g(x+11) = Kg((x;) = Σ decbi (a.bk(x;)cc (x:)ck,
(2)
where Å is eigenvalue associated with the eigenfunction du (x). which can be evaluated either the initial state Xy or any intermediate state x,, C₁ is the reconstruction coefficient, also known the "Koopman mode", which represents the projection of the observable function g onto the eigenspace. It immediately follows that g(xi+r) provided convenient and general framework to "predict and control" given dynamical system. Each Koopman mode evolves over time with its frequency and decay rate governed by the imaginary and real components, respectively.
20 XT
(x:xx(x))c (x,)ca for any т E N. This has
Koopman based techniques are particularly useful in data-driven setting because they only require measurements of observables. As such, they can be implemented even when the underlying model dynamics are unknown.
Dynamic Mode Decomposition (DMD). When using Koopman- based approaches, is critical to identify a suitable finite basis for representing the infinitediimensional Koopman operator: Dynamic Mode Decomposition (DMD) (Schmid) (010) one standard approach for inferring Koopman-bass models. It uses least-squares fitting techniques to approximate a finiteddimensional linear matrix operator, A, that advances high- high-dimensional measurements of a system forward in time:
g(x;-1)= R Ag(x;)
where A is an approximation of the Koopman operator, KC, in Eq. restricted to a measurement subspace spanned by direct measurements of the state x. Since the weight space of neural network is
OOr code is publicly available wttttttta//// http:::://gittbb.comm com/a wtt.:/cco////////////c////
3

---

Published as a conference paper at ICLR 2026
a fully observable system, we define g(.) to be the identity function in this work. That w; = g(w)). In practice, we often use "snapshots" " of the system arranged into two data matrices, W2 and W,+1. where columns of these matrices indicate measurements (i.e.. network weights) taken at certain time, and W++1 is W; shifted by one time step. Hence,
W1+1 ~ AW,
and A can be solved by A = W W W where W, UEVT is the Singular Value Decomposition (SVD), and W denotes the pseudo-inverse of W. A comprehensive discussion of DMD and its variants has been provided in Kutz et al. (2016).
DNN Training as a Dynamical System. There have been few works in recent years that adopt
Koopman-based approaches to accelerate the training process of a general-purpose DNN model
(Dogra & Redman 2020; Tano et al., 2020 Manojlovie et al.] 2020). (Dietrich et al. 2020) is generally considered the first work that establishes the connection between KOT and acceleration of numerical computation. Dogral (2020) is also one of the pioncer works but with a focus specifically on neural networks for solving differential equations. Generally speaking, these works take advantage of the predictive power of the KOT framework to directly predict network weights a few epochs later,
thus bypassing the consuming SGD iterations. However, we show in Fig. 2 that these methods tend to fail for larger network structures as the network size increases.
FGB lpr ser-eractive Pradictior TOT (Durs)
FGD Sor activi Predictor TOT (Durs)
Hon-seles ayn-selecive Precision
Fot (Oural
2.0
17an 1055
Tain Loss
10
1.2
16-
Epoch
Epoch
Epoch
(a) 2-layer FC
(b) 4-layer FC
(c) 6-1ayer FC
Figure 2: Performance comparison CIFAR-I0 using fully connected (FC) networks with varying depths, among SGD (iterative), PDT Koopman based predictive training where the predicted weights are applied to all parameters without checking the prediction quality (Tano et al.] 2020). Batch size=256, 1r-0.01. In our setup, for every three epochs of SGD, predictions are performed for the next five steps.
and the non-selective prediction, i.c.,
Directly predicting the evolution of neural network weights, by bypassing SGD, is inherently difficult due to the complex and unstable nature of training dynamics. The loss landscape is highly non-convex, filled with local minima, saddle points, and flat regions (Goodfellow et al., 2014), while the effective dynamics is non-stationary (Ghorbani et al. 2019). as both gradients and curvature shift as training progresses (Sagun al., 2017). In addition, neural systems can exhibit chaotic or highly sensitive regimes, where small perturbations quickly amplify and destabilize predictions (Li & Ravela [021]] Engelken et al., (023). This challenge is compounded by the stochastic noise introduced through mani-batch sampling.
Small prediction errors are highly sensitive and cumulative, risking divergence in the absence of continual gradient correction (Andrychowicz et al. 2016). Moreover, predictors trained in one context often fail to generalize across architectures and datasets, highlighting the difficulty of extracting universally valid patterns WWichrowska et al.) 2017: Metz et al. 2019). Together, these factors make weight prediction fundamentally unstable and error-prone task.
Beyond these general challenges, a number of prior works have attempted to predict future weights directly, such as Introspection (Sinha et al. 2017). WNN (Jang et al. 2023), and the more recent NiNo (Knyazev al. 2024). These methods typically rely on a separately learned predictor-either element- wise regression models or graph-basee networks- before being applied to new target model. The effectiveness therefore depends on the predictor's s meta- training distribution, and the inference cost grows with model size because predictions are applied at per-weight per-edge granularity. In contrast to these learned predictors, the proposed
on curated checkpoint datasets
4

---

Published as a conference paper at ICLR 2026
PDT adapts to the heterogeneous training dynamics of different parameters, enabling it to sustain network growth and provide a viable solution to the weight-predicet challenge. As a result, PDT functions lightweight plug-in without requiring external checkpoint datasets per-weight inference overhead. The efficiency of PDT has been validated on several benchmark models (e.g.. AlexNet, ResNet, and ViT), datasets (e.g., CIFAR-10 and ImageNet), spanning both supervised and self-supervised tasks.
METHODS
Let us first toy example demonstrate the effect of accelerating the learning of subset of variables to motivate the concept of differential learning. Consider the function.
fI, u. = x² + y² sin(2) + u²= cos(v) v) + w² ay + (2) + uow,
which involves six variables: 2, M/ u, v, To find the minimum of this function, employ simple GD optimization with a learning rate of 0.01. GD takes 53 steps to reach loss value below predefined threshold (0.1).
We then explore an alternative optimization strategy where the variables æ, y. use learning rate three times faster than that of the standard process, while u, U, optimized the normal rate but employing the updated values æ; y, 2. See Fig. X in Appendix A.. for the acceleration trajectory, where the trajectory maintains the same direction for and y but reaches the threshold in just 25 steps. This example demonstrates by strategically identifying a subset of variables and simply increasing their learning rate, the training can accelcrated about 53%. We also apply the proposed PDT to the same optimization problem and it reaches the threshold in 27 steps. See Fig. in Appendix A.] This toy example demonstrates the principle behind the idiom, rising tide lifts all boats!
3.1 PDT TRAINING FRAMEWORK
The PDT Training Framework addresses three key questions: 1) when to enable prediction, 2) how to integrate predictions with existing optimizers, and 3) which parameters should undergo accelerated updates. The complete PDT workflow and the mechanism involved in a single acceleration step illustrated in Fig. 3 The "prediction" block (P red)) is automatically but strategically placed among the baseline optimization blocks (OPT), acting as a plug-in enhancement within the existing optimization framework. Training begins with "Burn-in stage." where the model trained using the baseline optimizer for several epochs to accumulate sufficient history of weight snapshots. Following this stage, a prediction step performed with an adaptive interval, 7, to achieve accelerated learning.
Prediction
OPT
Pred
OPT
Pred
Last cosjectory
Neec 1 Step Ujprattes
Basidiee apriiiatiin
Combined Updones
Rediction
Lon-ttolltyyeee www.icecpppcccccccc.ooooo gits
Snopahota from last supocrs
Figure 3: Illustaaioon of the proposed PDT framework and the detailed mechanism for 7-step prediction, where qualified (or high-fidelity) predicted weights (red) and standard SGD-derived weights (blue) integrated that accelerate the training of the entire network.
5

---

Published as a conference paper at ICLR 2026
The bottom part of Fig. provides detailed illustration of how qualified predicted weights and standard SGD-derived weights are integrated together to achieve accelerated learning, as showcased in the toy example. The mask is governed by the dynamic consistency analysis to be elaborated in Sec. 3.2 If no element in the mask satisfies the criteria, then standard SGD-bbsed optimization takes over. The pseudocode of the complete PIJT algorithm is presented in Appendix A.2.2
The amount of computation required to perform a DMD-based prediction is comparable to that of a GD operation. It is important to note that the prediction operations are much less frequent (once for several epochs) compared to the standard GD operations (multiple times per epoch, depending on the batch size). Considering that PDT requires fewer epochs to reach convergence (see Table it can lead to significant computational savings and efficiency enhancements in the training of large-scale neural networks. The theoretical time and space complexity analysis is provided in Appendix AA.3 and detailed profiling and overhead analysis is provided in Appendix A.8
3.2 DYNAMIC CONSISTENCY ANALYSIS
The prediction step begins by applying DMD the weight snapshots. W. and W4+1, which yields a finiee-dmennsioal approximation of the Koopman operator, A, as in Eq 4 Since A is of high dimen- sion, N x N, where N is the number of weights (or parameters) of the network, is computationally intractable to solve directly. Hence, we resort to the Standard DMD algorithm (Tu et al. (014). This method projects the dynamics onto low-rank subspace to efficiently compute the eigenvalues and the high-dimensional DMD modes & without computing A directly (see Appendix A.2. for the derivation). To practically implement the spectral prediction from Eq. [ the predicted weight vector is computed from Eq. 5
prod w
A Φ ww;
(5)
where & is the matrix whose columns the DMD modes (approximating the eigenfunctions Φx), A is the diagonal matrix containing the corresponding eigenvalues λε. and की denotes the Moore- Penrose pseudoinverse. The term o (ii) projects the current state onto the DMD modes, calculating the Koopman mode amplitudes C₁ in Eq.
Our approach is based on the principle that DMD extracts the dominant patterns of the entire system's dynamics. However, at any given training stage, different parameters may exhibit varying degrees of alignment with these global patterns. Parameters experiencing rapid transitions, or local instabilities. may not conform to the global linear dynamics assumption underlying DMD. By leveraging the spectral components (Φ, A) derived from the low-rank approximation of A, we can perform multi-step prediction through more stable spectral evolution process using Eq. [S] which provides a prediction for the system's is global dynamics and also offers a perspective on the prediction for each parameter. The challenge, however, is how to determine whether such a prediction for each parameter has "high-fidelity" or "low-fidelity"
In fact, the correlation between the quality of prediction and training effectiveness has been heavily studied. From a ncuroscience perspective, the quality of predictions made by neurons is intricately linked to their learning dynamics (Schultz et al., 1997: Friston 2010). Accurate predictions lead to more stable and efficient learning, while poor predictions need stronger synaptic adjustments to improve future performance.
Therefore, we design a masking mechanism to identify parameters whose current local dynamics align with the system's is global dynamics, based the following two principles.
The acceleration effectiveness criterion. The absolute weight change between the predicted weight, pred and the current weight, w.rt at each epoch, i, should be larger than the absolute weight W change from the one-step optimization, wopt w.Pt. to enable accelerated learning; simultaneously, we impose an upper bound of 7 multiples of with - wopt where r is the prediction step, to ensure stable convergence. That is,
(wwpp opt ++1
opt
www.eee prod
S 7/|
opt #+
w Wott
This criterion ensures that the prediction provides significant advancement beyond what single-step optimization would achieve, making the acceleration worthwhile. See Appendix A.4 for convergence guarantee analysis.

---

Published as a conference paper at ICLR 2026
The dynamic consistency criterion. The direction of weight change from prediction should align with the local gradient- based dynamics. That is, the temporal evolution captured by the global DMD analysis should be consistent with the current local optimization trajectory. Specifically:
pred pred w+k--1,) sign(wi+t) H+1j w.)) fign(W+tkJ
(7)
where j is the index for each element in the weight vector and k {1, {1.....7)) Note that when
.7).
I, w.rred wiph This criterion ensures that each step of the prediction interval follows the J same trend of growth as that of the local optimization.
Based on these two principles, a mask, m can be constructed with its element equal if both Gqs. 7 are satisfied; otherwise, the corresponding element is zero. This dynamic consistency analysis evaluates these two criteria independently for each parameter. Parameters satisfying both criteria are deemed to be in predictable evolutionary phase, allowing safe application of temporal acceleration through the global dynamic model. Parameters failing these criteria may be experiencing complex local behaviors (such rapid transitions, oscillations, instabilities) that deviate from the global linear dynamics assumption, requiring fallback to gradient-based updates. Note that Eq, rigid criterion to enforce not only that the final predicted weight changes align with the local optimization direction, but also that every intermediate step in the predicted trajectory maintains direction consistency.
4 EXPERIMENTS
4.1 GENERALIZATION STUDY OF PDT
We evaluate the effectiveness of the proposed PDT framework in accelerating learning across variety of popular neural network architectures with different scales, including Delly-Convolutioaa--Neo -Ntwoork (FCN) (3.9M parameters), AlexNet (57M parameters), ResNet-50 (25.6M parameters), ViT-Base (86.4M 46.4M parameters), and ViT-Huge (632M parameters). PDT shows significant advancement over previous prediction-only Koopman-bans based methods limited to simpler models [e.g.. (Tano al. 2020) with 2.9M trainable parameters). We also use range of optimizers, including SGD, SGD with momentum, and A (Loshchilov Hutter 2019). Our validation spans multiple benchmark datasets, from CIFAR-10 to the more challenging ImageNet-IK.
0.025
107
Achieved
0.000
8.7
1014
0.70
Smo 3mtn
ON fister
Train
50D
50D
0.050
PDT
Basclin
Sest
025
200
420
202
420
700
tseconds)
(peconds)
Figure 4: Analysis of PDT's Optimization and Generalization Efficiency. Trained on CIFAR-10 using AlexNet, batch size=256, 1r=0.05, with Cosine Annealinnn scheduler.
Fig. 4 iillustrates the efficiency of PDT by showing the training curve on CIFAR- 10 using AlexNet. We assess PDT's performance from two perspectives: (1) Optimization Efficiency, measured by the Time to Baseline Best Train Loss (or TTB-Loss), quantifies the speed at which PDT converges on the training objective. (2) Generalization Efficiency, measured by the Time Baseline Best Validation Accuracy (or TTB-Ace,, reflects the speed at which PDT obtains useful, well-generalized model. Note that all runtime metrics reported in this paper represent the total wall-clock time, fully inclusive of all computational overheads introduced by PDT, such SVD decomposition, multi-step prediction, and mask generation. The detailed training curves on various network structures can be found in Fig. 10 in Appendix A.5 A detailed profiling of these overheads is provided in Appendix A..8
In all experiments, unless otherwise specified, we use the past five epochs to form the snapshot with one-epoch interval to predict weights in the next five steps. Prediction starts from the 5th epoch.
7

---

Published as a conference paper at ICLR 2026
A comprehensive study of PDT's performance under different training configurations (batch sizes, learning rates, and optimizers) and hyperparameters [i.e., prediction steps (T). prediction interval (I)), starting epoch (To), and past snapshot counts (h)] can be found in Appendix A.6 demonstrating robust performance across various training hyperparameters.
As elaborated in Sec. the computational load of the Koopman-related calculations is comparable to that of batch-level updates. However, since we apply these calculations at the epoch level, the overhead introduced by DMD is effectively compensated by the acceleration in loss reduction. We observe from both Table and Fig. TO that the proposed PDT consistently achieves the best training loss of the Baseline but in fewer number of epochs without sacrificing performance. All experiments were repeated with five random seeds to ensure reliability. Unless otherwise specified, all results in the tables are reported as mean + standard deviation over five runs.
Table 1: Runtime comparison. FCN and AlexNet are trained on a single Nvidia RTX A6000 GPU, while ResNet-50, ViT-Base, and ViT-Huge are trained on three Nvidia H100 (80 GB) GPUs. Using the same experimental setup and hyperparameter configurations as in Fig. TO
Baseline Optimizer
TTB-Loss (s)
TTB-Ace (s)
Runtime Reduction (%)
Model
Baseline
PDT
Baseline
PDT
Train Loss
Val. Acc.
FCN
SGD
2174.32
1313.52
2088.58
1424.14
39.59
31.81
AlexNet
SGD
683.93
430.91
531.30
347.11
37.00
34.67
ResNet-50
SGD-M
110063.72
88752.33
121449.60
92133.34
19.36
24.14
ViT-Base
A
259241.21
232810.62
296028.36
243097.58
10.20
17.88
ViT-Huge
Adam W
725564.86
653854.05
741220.54
660711.80
9.88
10.86
The last column in Fig. I0 illustrates so-called "masked ratio curve" unique to PDT, where it tracks the percentage of predictions accepted according to the masking strategy described in Sec. 3.2 We make two interesting observations. First, observe that the masked ratio always starts with higher values in the early stage of the training process, then generally decreases as training progresses. This trend aligns naturally with how GD-based baselines behave. That is, in the early stage of the training process, the loss landscape is typically easier to optimize. leading to faster reduction in loss and more stable gradient directions, which in turn allows more weights to pass the masking criteria, hence higher masking ratio. Later in training, gradients tend to oscillate more around the optimum, making more challenging to predict, hence less percentage of predicted weights being accepted. More interestingly, we observe that smaller networks on simpler tasks (FCN/AlexNet on CIFAR-10) show relatively more gradual reduction in the masked ratio, while larger networks more complex tasks (ResNet-S0/ViT on ImageNet) exhibit a much sharper reduction of masked ratio, especially at the early stage of the training process. This pattern implies that for large networks large datasets, the training dynamics are inherently complex and challenging to predict. The higher masked ratio at the initial training stage is primarily attributed to the steep loss landscape and large gradients. As the model converges and gradients diminish, the intrinsic complexity of the training dynamics becomes dominant, resulting in a rapid reduction in the percentage of weights that can be convincingly predicted (according to the proposed masking strategy). A detailed analysis of how the mask distribution evolves across different layers of the network during training is provided in Appendix A.7 A study on the impact of different initial learning rates on the masked ratio evolution is presented in Appendix A.
the GD>based optimizer approaches (local) minima,
To further demonstrate the versatility of PDT as a plug-in enhancement, we extended our evaluation to include a broader range of optimizers [SGD, SGD with momentum, Adam, AdamW, RMSprop, Shampoo (Ciupta et al., 2018), and LAMB (You et al. [0200] while keeping the network architecture and other configurations fixed. The results in Table 2 show that PDT consistently reduces the time to reach baseline best loss across these optimizers.
We also extend our evaluation to the domain of self-supervised learning (SSL). We select Sim- Siam (Chen & He 2021). prominent con-contrastive method, as our testbed. SimSiam's 's training dynamics, which are driven by a stop-gradient mechanism and negative cosine similarity objective. are fundamentally different from those of supervised learning. The results, summarized in Table demonstrate that PDT's advantages generalize effectively to SSL.
8

---

Published as a conference paper at ICLR 2026
Table Impact of baseline optimizers PDT performance. Trained on CIFAR-10 using AlexNet, batch size=256, momentum=0.9, with Cosine/mmealingLL scheduler.
Runtime Reduction (%)
Optimizer
Ir
Final Accuracy
Best Train Loss
TTB-Loss
SGD
0.7930 = 0.0023
0.0002 1 0.0000
665.27 1 9.08
0.1
19.67
PDT
0.7978 = 0.0032
0.0002 1 0.0000
534.41 1 12.64
Momentum
0.6672 = 0.0068
0.8609 1 0.0166
752.74 1: 9.62
0.001
41.06
PDT
0.7298
0.0051
0.5358 + 0.0165
443.68 + 8.75
Adam
0.7952 = 0.0063
0.0001 + 0.0000
779.13 + 11.81
0.0005
14.87
PDT
0.8050 0.0050
0.0002 + 0.0000
663.28 + 15.30
AdamW
0,8031 = 0.0021
0.0077 + 0.0002
652.92 + 5.26
5e-5
28.36
PDT
0.8149 0.0037
0.0013 + 0.0002
467.76 + 6.05
RMSprop PDT
0.7996
0.0032
4.1e-5 + 1.6e-5
661.08 + 0.42
0.0001
15.35
0.8108 0.0026
2.4e-5 + 0.6e-5
559.61 0.00
Shampoo PDT
0,8012 = 0.0071
0.0040 + 0.0005
736.56 1 10.58
0.001
16.03
0.8101 0.0043
0.0033 1 0.0003
618.49 1 12.72
LAMB
0.8034
0.0025
0.1215 1 0.0085
663.25 1 5.44
0.001
44.24
PDT
0.8140 0.0027
0.0036 1 0.0006
369.82 1 10.59
Table 3: Performance comparison of SimSiam pre-training on CIFAR-I0 with a ResNet-18 18 backbone. trained for 200 epochs, 1r=0.03, batch size=256, momentum=0.9, with CosineAnmealingLR scheduler.
Runtime Reduction (%)
Optimizer
Final Accuracy
TTB-Loss (s)
TTB-Acc (s)
Train Loss
Val. Acc.
SGD Momentum
0.7285 1 0.0166
9611.35 1 837.98
7353. 1253.12-4 1 1063,96 6169.04 1 1142.97
48.78
16.10
PDT
0.7685 1 0.0144
4922.92 712.55
In addition, further analysis of PDT's performance under non-i.i.d. training conditions is presented in Appendix A.11 The cross-domain evaluation on natural language processing tasks is in Ap-
pendix A.12
4.2 MASKING STRATEGY: RANDOM SELECTION AND VALIDATION Loss?
We study the effectiveness of the proposed masking strategy by comparing it with two ad hoc strategies. First, we perform series of runs in which we randomly select subsets of weights and increase their learning rates. In Fig. 5 regions highlighted in green show results from different runs where subsets of weights had their learning rates increased to match the step number used in the predictions. The selection ratio used here matches the average masking ratio applied during PDT. The results indicate that randomly accelerating weights cannot match the performance improvements from PDT and often lead to significant instability during training.
0.7
NARI4A 20 LON4
3.5
0.5
Tran
Velideelon
0
Prodictoo
0.1
Epoch
Epoch
Epoch
Figure 5: PDT vs. randomly selected subsets with higher learning rates (with the same mask ratio). Trained on CIFAR-10 using AlexNet, batch size=256, Ir 0.05.

---

Published as a conference paper at ICLR 2026
Second, we perform a series of rums where subsets of Koopman predicted weights randomly selected. The regions highlighted in green in Fig. show the outcomes of these trials. Quite frequently, these rums result in gradient explosions, leading to non-recoverable errors (NaN values) in subsequent epochs. This experiment underscores the importance of a principled masking strategy in Koopman Training. Random masking, without considering the training dynamics, can lead severe divergence and training failure. Our findings highlight that strategic selection based on "high-fidelity" predictions is crucial the success of PDT. Furthermore, conduct comprehensive ablation study in Appendix A.9 to validate the necessity of combining both the acceleration effectiveness (Eq. and dynamic consistency criteria (Eq.)
SGD
2.50
2.50
Reciation
2.2
LOS5
Val ldttto ACCURSC
Vay ldvtto
2.75
1.50
1.2
525
50D
Racker
Rander
Eooch
Figure 6: PDT random mask prediction (with the same mask ratio). Trained CIFAR-10 using AlexNet, batch size=256,, =r 0.01.
We implement another baseline scheduling scheme that switches between prediction and SGD based on the validation loss trend: apply prediction when validation loss decreases and roll back to SGI) updates when validation loss starts increase. Fig. iillustrates the training dynamics under this strategy. Initially, DMD is engaged due to its slight advantage in reducing validation loss. However, as training progresses, significant surge loss observed, suggesting misalignment between the DMD-pre prediccdd weights and the optimal trajectory for the network. Even after reverting to SGD, the model failed to recover, indicating that relying solely on validation loss as a trigger for switching between PDT and SGD is inadequate.
Figure 7: Performance comparison on CIFAR-10 using AlexNet: SGD vs. Koopman- prediction (switching between prediction and SGD based validation loss). L: Train loss. R: Validation loss.
5 CONCLUSION
This paper proposed a novel predictive differential training (PDT) framework based on the study of training dynamics. PDT incorporate the idea of "differential learning" into the predictive training framework for accelerated learning even for complex network structures. The key contribution is the design of effective masking strategy based on dynamic consistency analysis, which selects only those predicted weights of high-fidelity whose local training dynamics align with the global dynamics. Analogous to the saying a rising tide lifts all boats, our setting, a subset of high-fidelity predicted weights facilitates more efficient training across the entire network!
The training process deep network with millions to billions of parameters indeed presents an intriguing dynamical system that the control community has not faced before. This would stimulate further investigation into the development of better data-driven dynamical system analysis algorithms in addition to DMD. Innovative approaches, such streaming DMD (Hemati et al. 2014 Liew et al.) 2002), can not only reduce the memory footprint of constructing trajectory matrices, but also improve computational efficiency.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
portion of the computation for this work was performed on the University of Tennessee Infrastruc- ture for Scientific Applications and Advanced Computing (ISAAC) computational resources. We also thank the anonymous reviewers for their constructive comments that significantly improved this paper.
REFERENCES
Marein Andrychowiez, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Advances neural information processing systems, 29, 2016.
Gal Berkooz, Philip Holmes, and John L Lumley. The proper orthogonal decomposition in the analysis of turbulent flows. Annual review of fluid mechanics, 25(1):539-575, 1993.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 15750-15758, 2021.
Jacob Devlin, Ming-Wos Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv:1810.04805. 2019.
F. Dietrich, T. N. Thiem, and I. G. Kevrekidis. On the koopman operator of algorithms. SIAM Journal on Applied Dynamical Systems, 19(2):660888,, 2020.
Akshunna S Dogra. Dynamical systems and neural networks. arXiv preprint arXiv:2004. 1826, 2020.
Akshunna S Dogra and William Redman. Optimizing neural networks via koopman operator theory. Advances in Neural Information Processing Systems, 33:2087-2097, 2020.
John Duchi, Elad Hazan, and Yoram Singer. Adagrad: Adaptive subgradient methods for online
learning and stochastic optimization. Journal of Machine Learning Research, 12:2121-2159, 20
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 201 lb.
Rainer Engelken, Fred Wolf, and Larry F Abbott. Lyapunov spectra of chaotic recurrent neural networks. Physical Review Research, 5(4):043044, 2023.
K. Friston. The free-energy principle: unified brain theory? Nature Reviews Neuroscience, 11(2): 127-138, 2010.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning, pp. 2232-2241.
PMLR, 2019.
Ian J Goodfellow, Oriol Vinyals, and Andrew M Saxe. Qualitatively characterizing neural network optimization problems. arXiv preprint 4rX1v:4412.6544, 2014.
Antonio Gulli. AG News Dataset. httt::/grcops..dd un
http::/rrroo..dd.........//////p http://.......ggll//...oooo. Accessed: 2025-12-01.
of ww_.ews_arrtte wewiiielss... html
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- tion. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, and Mu Li. Bag of tricks for image classification with convolutional neural networks. arXiv: 187, 2019.
Maziar S Hemati. Matthew O Williams, and Clarence W Rowley. Dynamic mode decomposition for large and streaming datasets. Physics of Fluids, 26(11), 11), 2014.

---

Published as a conference paper at ICLR 2026
Juhwan Jang, Dongyoon Kim, Kihyuk Kim, Minsu Cho, and Bohyung Han. Learning to boost training by periodic nowcasting near future weights. In International Conference on Machine Learning, 2023.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6900 2014.
Boris Knyazev et al. Accelerating training with neuron interaction and nowcasting networks (nino). arXiv preprini arXiv:2401.12344 2024.
J. Nathan Kutz, Steven L. Brunton, Bingni W. Brunton, and Joshua L. Proctor. Dynamic Mode Decomposition: Data-Driven Modeling of Complex Systems. SIAM-Society for Industrial and Applied Mathematics, 2016. ISBN 778--611197-44--.
Mu Li. David G Andersen, Jun Woo Park, Alexander Smola. Amr Ahmed, Vanja Josifovski. James Long, Eugene Shekita, and Bor- Yiing Su. Scaling distributed machine learning with the parameter server. In /th [USENIX] Symposium on Operating Systems Design and Implementation ((OSDI) 14), pp. 383-598. 2014.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, al. Pytorch distributed: Experiences on accelerating data parallel training. arXiv preprint arXiv:2006. 2020.
Ziwei Li and Sai Ravela. Neural networks geometric chaotic maps. IEEE Transactions on Neural Networks and Learming Systems, 34(1))527-533, 2021.
Jaime Liew, Tuhfe Göşmen, Wai Hou Lio, and Gunner Chr Larsen. Streaming dynamic mode decomposition for short-term forecasting in wind farms. Wind Energy, 25(4):719-734, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer-
ence on Learning Representations, 2019. URL
Bkg6RiCqY7
Iva Manojlovic, Maria Fonoberova, Ryan Mohr, Aleksandr Andrejcuk, Zlatko Drmac, Yannis Kevrekidis, and Igor Mezić. Applications of koopman mode analysis to neural networks. arXiv preprint arXiv:2006. 11765, 2020.
James Martens. Deep learning via hessian-free optimization. In ICML, pp. 735-742, 08 2010.
Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Understanding and correcting pathologies in the training learned optimizers. In International Conference on Machine Learning, pp. 4556 4565. PMLR, 2019.
I. Mezić. Spectrum of the Koopman operator, spectral expansions in functional spaces, and state-space geometry. Journal of Nonlinear Science, 30(5):2091-2145; 2020.
I. Mezić and A. Banaszuk. Comparison of systems with complex behavior, Physica Nonlinear Phenomena, 197(1-2):001-1333 2004.
Igor Mezić, Spectral properties dynamical systems, model reduction and decompositions. Nonlin- ear Dynamics, 41:309-325, 2005.
Igor Mezić. Koopman operator, geometry, and learning of dynamical systems. Notices of the American Mathematical Suciety, 68(7):1087-1105. 2021.
William Redman, Maria Fonoberova, Ryan Mohr, Ioannis Kevrekidis, and Igor Mezić. An operator theoretic view on pruning deep neural networks. International Conference on Learning Representations (ICLR), 2022.
William Redman, Juan Bello- Rivas, Maria Fonoberova, Ryan Mohr, Yannis Kevrekidis, and Igor Mezic. Identifying equivalent training dynamics. Advances in Neural Information Processing Systems, 37:23603-23329, 2024.
12

---

Published as a conference paper at ICLR 2026
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathematical Statistics, 22(3):400 407, 1951. doi: 10.1214/amm//1177725666 URL https: 10 10.12244/omm/1177/
David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-propagating errors. nature, 323(6088):333-536, 1986.
Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over- peee-paraetrizedd neural networks. arXiv preprint vrXiv:1706.04454 2017.
Peter Schmid. Dynamic mode decomposition of numerical and experimental data. Journal of Fluid Mechanics, 656:5-28. 2010. doi: 10.1017/500221 000000022122211111 12010001217.
W. Schultz, P. Dayan, and P.R. Montague. A neural substrate of prediction and reward. Science, 275 (5306):1533-1599 1997.
Abhishek Sinha, Mausoom Sarkar, Aahitagni Mukherjee, and Balaji Krishnamurthy. Intro- spection: Accelerating neural network training by learning weight evolution. arXiv preprint arXiv:/704.04959) 2017.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International Conference on Machine Learning (ICML), pp- 1147. PMLR, 2013.
Mauricio E Tano, Gavin D Portwood, and Jean C Ragusa. Accelerating training in artificial neural networks with dynamic mode decomposition. arXiv preprint arXiv:2006. 14371, 2020.
Tijmen Tieleman and Geoffrey Hinton. Rmsprop: Divide the gradient by running average of its recent magnitude. coursera: Neural networks for machine learning. COURSERA Neural Networks Mach. Learn, 17, 2012.
Jonathan H. Tu, Clarence W. Rowley, Dirk M. Luchtenburg, Steven L. Brunton, and J. Nathan Kutz. On dynamic mode decomposition: Theory and applications. Journal of Computational Dynamics, 1(2):391 421, December 2014. ISSN 2158-2491. doi: 10. 10.3934/jjd URL https:///// www a www.mmssiineeee org/en.////..oo ottp://nnaarrtccee//o//////33//3 1.91. Publisher:
Journal of Computational Dynamics.
Fanqi Wang, Weisheng Tang, Maofeng Tang, Konstantinos Georgiou, Hairong Qi, Cody Champion, and Mare Bosch. Koopman- based transition detection in satellite imagery: Unveiling construction phase dynamics through material histogram analysis. In IGARSS 2024-2024 IEEE International Geoscience and Remote Sensing Symposium, pp. 8786-8790, IEBE, 2024.
Olga Wichrowska, Niru Maheswaranathan, Matthew W Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Nando Freitas, and Jascha Sohl-Dickstein. Learned optimizers that scale and generalize. In International conference on machine learning, pp. 3751 PMLR, 2017.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song. James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2020. URL httttp:://openzeieew net hett:/orrumm..--yyyxvv...e.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification, Advances in neural information processing systems, 28, 2015.
13

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 CONVERGENCE PATH OF THE TOY EXAMPLE
To better illustrate the effectiveness of differential learning strategies, we designed a toy optimization problem with six variables. See Sec. 3.1 for the description of the problem. Starting from the initial point [2.0, 2.0, 1.0, 0.5, - 0.5, 1.5] with learning rate of 0.01. Fig. 8 shows the optimization trajectories in the x-y plane, where the background color represents the function value at each point. The blue line with dots represents the GD trajectory, while the red dashed line shows the path of accelerated GD where x, y. Z variables use 3x learning rate. All points on the trajectories represent the state after each optimization step. The arrows indicate where each method reaches the threshold value (0.1). Building upon this observation, we apply our proposed PDT method to the same optimization problem. Fig. presents the comparison between standard GD and proposed PDT method the same toy optimization problem. Fig. (aa uses the same visualization scheme as Fig. g showing how PDT follows similar path but reaches the threshold faster (PDT reaches the threshold in 27 steps).
Fig. ()) clearly demonstrates the acceleration effect, where PDT's loss decreases more rapidly than
GD. The horizontal dashed line indicates the threshold value used as the stopping criterion.
306
24-
27.0
26-
214
15
110
12.
nacted theshes
32.6
Aocela sstd resct Fveshes
-15-
Figure 8: The differential learning trajectory of the toy example provided Sec. 3.1 y dimensions are shown.
Only the x and
306
16-
PS1
25
370
214
P0Trucee Foctee di
-45
31
110
-25
2.0
25
2.0
(b)
Figure 9: Performance comparison between GD (53 steps) and PDT (27 steps) on the toy optimization
problem. (a) Optimization trajectories in the x-y plane. (b) Loss values during optimization.
14

---

Published as a conference paper at ICLR 2026
A.2 IMPLEMENTATION DETAILS OF PDT
A.2.1 DERIVATION OF THE DMD ALGORITHM
In the main text, we introduced the Koopman operator approximation A W.'W. However, as noted, an N N matrix (where N is the number of parameters, typically millions or billions), making direct computation intractable. In our implementation (see Appendix [A.2.2),
employ the DMD algorithm (Tu et al. 2014) to compute the spectral decomposition of A in a low-rank subspace. This section provides the step-by-ster derivation linking the snapshot matrices W,, W4+1 in Eq 4 to the spectral components Φ A used in Eq. 5
Let w, E RN be the flattened weight vector of the neural network at the current epoch i. To capture the training dynamics, we construct two snapshot matrices using the weight trajectories from the past epochs: Input Matrix (corresponding to W; in Eq. 4 contains the sequence of weights from the history buffer, excluding the last weight state; Shifted Matrix (corresponding to W4+1 in Eq, 4 contains the same sequence shifted forward by one time step, ending with the current weight w, For the remainder of this derivation, we refer to these matrices as W; and W4+1 to maintain consistency with the main text. Both matrices are in [Rxx(h1) where h < (e.g., h 12=55) 5).
We first compute the reduced SVD of W;:
W UEVI
(8)
where U E RNr, E R and V 1)x-) Here S h is the truncation rank. This step reduces the dimensionality from the vast parameter space to the small snapshot space r.
Instead of computing A = W w = W++1 V we project the high-dimensional operator A onto the low-dimensional subspace spanned by the proper orthogonal decomposition (POD) modes U (Berkooz al. 1993). We compute the proxy matrix À Rrxr:
(9) Computationally, this involves multiplying the large matrix W++1 by small matrices V and Σ-1, then projecting onto U. This results in tiny matrix that captures the essential dynamics of the full system.
A= AU = U (W.+ Y = U W
Since À is small (r x r), we can efficiently compute its eigendecomposition:
Āv VA
(10)
where A = diag(Ai,..... .,,,,..,.pp) contains the eigenvalues (which approximate the eigenvalues of the full operator A), and v E Crxx contains the eigenvectors of Á.
The eigenvectors L are in the low-dimensional subspace. To map the eigenvectors back to the full parameter space, we use the Standard DMD) formulation, which computationally efficient and numerically stable:
U
(11)
This yields the DMD modes CNxr, This step is crucial as it provides the mapping basis for our prediction equation. (Note: The Exact DMI) formulation +- W+1VV--1 is also option, but incurs additional computational cost).
With & and A computed, predict the future state 7 steps ahead, starting from the current weight w;. The prediction equation (matching Eq, in the main text) is derived as:
pred ++
Re( w:)
(12)
Here, the term c = o w, represents the projection of the current weights onto the DMD modes (i.e., finding the mode amplitudes, which corresponds to the coefficients с₁ in Eq. 2 Since the modes + are generally non-ont non-orthogonal (as A is not symmetric). & denotes the Moore-Penrose pseudoinverse, which provides the least-squares solution arg mina |w; -6c/2. Finally, we take the real part of the result, as the neural network weights must be real-valued.
A.2.2 ALGORITHM PSEUDOCODE
15

---

Published as a conference paper at ICLR 2026
Algorithm PDT algorithm
Require: baseline optimizer Obase past snapshot counts h, start epoch for prediction To. predicted steps T. prediction interval T, number parameters N
Ensure: Trained model parameters
1: Initialize weight history matrix Wixhs counter Ce
2: for epoch i
to T do
3: if 2 To and Ce 2 T; then
4:
Obtain wopt (i 1) from history matrix W xh
5:
Train model for one epoch using Obase, save weights after training as w
6:
Calculate DMD from W Nx obtain modes & and eigenvalues A
7:
1) Decompose: + "P (t) (Project current state onto modes)
8:
Evolve:
Compute future state w-Prod (i + r) f- Re( Re(
9:
3) Masking: Create mask M by comparing dynamics (Eqs. and
SGD step: Asgd = ( - w
PDT step: Apdt w (i 7) (
12:
4) Assemble: Update weights selectively
13:
w(i) + o wpred (i r) (1 M) @ woPe(1)
14:
Update model parameters with updated w(i)
15:
Ce 00
16:
else
17:
Train model for one epoch using Obase
18:
Ce Ce
19:
end if
20:
Update weight history matrix W xxh
21:
end for
A.3 COMPUTATIONAL COMPLEXITY ANALYSIS
To provide rigorous understanding of PDT's efficiency, we analyze its complexity in terms of both computation time and memory usage. We consider DNN with N parameters trained on a dataset with 5 samples.
Time Complexity. The computational load for processing each batch using standard SGD is directly proportional to both the batch size (B) and the number of parameters (N), resulting in complexity of Ó( N) per batch, or O(S N) per epoch.
Integrating Koopman operator predictions into the DNN training process entails constructing snapshot matrix from h past epochs of the parameter trajectories, with the matrix dimension being N h. As derived explicitly in Appendix A.2.I the prediction process involves several steps. The dominant operations and their complexities are:
SVD of W;: O(N x h?))
Computing A: O(N h)2)
Eigendecomposition of Á: (hh))
Computing Modes Φ: O(N x 122)
Prediction (Solve Φ(): O(N h²)
The total complexity is O(N x 12). Since h is small constant (e.g., h 5) in our experiments while N can reach millions or even billions, the quadratic impact of h remains manageable relative to N. Since the prediction step occurs only once per epoch (or every few epochs, depending T)), the amortized cost is minimal compared to the O(S x N) cost of the baseline optimization over the full dataset.
Space Complexity. The additional memory requirement for PDT dominated by the storage of the weight history matrix W, which stacks h snapshots of the model parameters. Thus, the space complexity is ONN x h). For our default setting of h = 5, this corresponds to storing additional copies of the model weights. On modern training hardware (e.g., NVIDIA A100 with 40GB+ VRAM),
16

---

Published as a conference paper at ICLR 2026
this overhead is manageable. For instance, ResNet-50 model (N 25.5M) requires approximately 100 MB per snapshot (in float32 precision), totaling ~ 500 MB for 5, which is minor compared to the memory consumed by activation maps and optimizer states. Note that these snapshots stored only temporarily and are overwrittnn after each DMD computation, so the space cost does not accumulate over training-
detailed runtime and memory profiling analysis is provided in Appendix A.8
A.4 CONVERGENCE ANALYSIS
We analyze the convergence of the hybrid update that combines DMMDbased predictions with gradient- descent updates under the masking strategy in Sec. 3.2 Throughout this section, we use wy to denote the current parameters and [(w) denote the loss.
A.4.1 ASSUMPTIONS AND UPPER-BOUND CONSTRAINT
Assume that L: LRR R has L-Lipschitz continuous gradients, i.e., [Vv(())
V _ L|| iz - yll.
Vx,y R
Let wint denote the parameters after one step of the baseline optimizer
opc ++1
w opt 1/ ((w; (Pt),
with step size O. Let w prod denote the step DMD prediction starting from wipe. As Eq in
i+t
Sec. 3.2 PDT enforces the magnitude bound
|
w TM VL(w;
(13)
In addition, for each element j accepted by the mask, the dynamic consistency criterion (Eq. ensures the sign alignment, sign(w pred wrod Hence, we have
(4k-1.)) sign(w opt
w.y") where k 1, 1,
sign(wprod prod
w opt = sign(w (+11j (6 = sign(-n) V ((w
This leads to
(wpted
20 (VL(((()) O.
A.4.2 UPDATE RULE WITH MASKING
Let m, E {0, 1 denote the binary mask. PDT forms the next iterate by opt
pred W,+1 O w i+
+ (1 m.) C Wi41
(14)
where : denotes elementwise multiplication. Define
= Wi+1 oot
Elementwise,
pred HTi
3 mij
opt
di.j
(15)
17 VL
mij =
For masked elements, [dر,jl puppeed |d,,| = v/ VL(W S Tyl [(w
1 $ Tn VL opt) by Eq. 13 For unmasked ones,
opt
since r 2 1. Hence |dا S TT) V ( for all j. and
||,ll| - tn WL(w |||
(16)
Furthermore, for elements where m,j = 1, the acceleration effectiveness criterion (Eq. o lower bound) implies |dayl > Vl Combined with the sign consistency,
(( da,s -n (V
for all j. Summing,
WL(W,"t) ), - [n
(17)
17

---

Published as a conference paper at ICLR 2026
A.4.3 DESCENT INEQUALITY AND CONVERGENCE
Because C has L-Lipschitz gradients, the standard descent lemma yields
L
C(x +d) d) [(x) + (VL(x), d) |
(18)
Applying Eq. 18 with x = wit and d = di and using Eq. 16
17
we obtain
[(w++1) - L(w 7l W. )1 + 2
WW opt
(19)
2
Grouping the terms.
[(wi+1) [(w,pt) 11 ( LIfT²
)
/
(20)
2
sufficient condition for the coefficient,
Lyr² 2
be positive
2
1) 5 L
(21)
Under Eq. 21 we obtain the GD-style descent inequality
[(w:+1) < [(w.pt).
(22)
Thus (C(w,)} } monotonically non-increasing and bounded below, hence convergent. Standard results for gradient-based methods with bounded steps and L-smooth losses (e.g., inexact gradient descent) then imply that
||V[(w;)] +0
and every limit point [w;} is stationary.
A.5 DETAILED TRAINING CURVES ON VARIOUS NETWORK STRUCTURES
Fig. T0 presents the detailed training curves and performance comparison of PDT and baseline optimizer across various network structures.
A.6 EFFECT OF TRAINING HYPERPARAMETERS
Several primary hyperparameter' require careful consideration in PDT:
Prediction Steps (T): Derived from DMD, the number of prediction steps significantly influences the training speed. As shown in Fig. [1 a). training accelerates within certain range of prediction steps. However, extending beyond a critical threshold, such as nine steps in our study, can introduce large errors and potentially cause gradient explosion.
Prediction Interval (T)): The interval between Prediction blocks impacts the effectiveness of acceleration, depicted in Fig. [T).. shorter interval can enhance training speed if the predictions are accurate. Nevertheless, the quality of predictions may decline as the training progresses, rendering the network more sensitive to errors, particularly as nears convergence.
Starting Epoch (To): The starting epoch for acceleration must be greater than or equal to the number of epochs used to build the snapshot, as illustrated Fig. ITc).. The initiation of acceleration is influenced by factors such as initialization, learning rate, and model architecture.
Past Snapshot Counts (h:) Fig. d) indicates that the number of epochs needed to construct the snapshot matrix for prediction also influences the train loss. This value cannot be too small or too large. If is too small, the snapshot will not have sufficient measurements to precisely estimate the dynamics of the training process. If it is too large, DMD would have missed the local dynamics with only a coarser grasp the general training dynamics.
Overall, these PDT-related hyperparameters robust across optimizers, architectures, and datasets. Based on our experience, below are "Rule of Thumb" guidelines to help find the appropriate hyperpa- rameters for different scenarios.
18

---

Published as a conference paper at ICLR 2026
107
SGD
0.33
0.30
Looe
0.6
0.25
20-1
0.-
0.20
Toin Lock
0.4
0.15
0.3
0.10
40"4
Epoch
Eooch
(a) Trained on CIFAR- 10 using FCN, batch size=256, 1=0001, with CosineAnnealingL.. scheduler.
109
2.5-
Beer
Raalls
2.0
Valcaatiin
Than Loss
2.5
0.10
60-44
10a2
0.00
Eooch
Epoch
Epach
Eboch
(b) Trained on CIFAR-10 using AlexNet, batch size=256, Ir=0.05, with CosincAnncalingL... scheduler.
0-337
P3T
S2D with Momentum
SGD with Nomertum
PDT
0.35
ror
Seat
Veldstin Accuety
0.30
Valiacioo LGG
65
Meseed Rallo
0:25
54
0.20
0.35
Tain
0.30
S00
0.05
00
PDT
0.00
1.00
250
100
120
200
302
100
202
300
Epoch
Epoch
Epoch
Epoch
(c) Trained ImageNet- IK using ResNett-O0, batch size=1800, 1r=0.1, momentum=0.9, with CosineAn-
nealingLR scheduler.
PDT
Adsmii
Acantt HDT
Adammi? P3T
0.25
PCT
Train scak)
Baseline Lens
Valldetion Loos
Vallcotton ACCursYy
0.00
0,13
0.30
0.2
0.05
6.00
200
400
Epoch
Epoch
Epoch
Epoch
(d) Trained ImageNet- KK using ViT-Base, batch sizc=1800, Ir=00003, momentum=0.9, with Cosine nealingLR scheduler.
Figure 10: Performance comparison between baseline optimization and PDT.
Past Snapshot counts (h): is "sweet spot". Smaller h is insufficient for capturing dynamics, larger includes "stale" weights from much earlier training and introduces additional overhead. For networks with extremely dynamic changes, larger value of (e.g., /=100) is also worth trying.
Prediction Steps (T): We recommend starting with 5 as robust default value. This value provides a good balance between acceleration benefit and prediction accuracy across diverse architectures and datasets. Users can increase 7 to 7 if their training loss curves very stable and exhibit minimal variance. Conversely, if gradient explosion or instability occurs, reducing to 3 provides more conservative acceleration while maintaining stability. A practical configuration is to set 7 E [3, 7]. Too large leads to divergence, while too small makes it meaningless.
19

---

Published as a conference paper at ICLR 2026
Prediction Interval (T;): We recommend setting T; = 1 as the default. If the training process is unstable, then gradually increase the interval.
Start Epoch (To): The start epoch To should typically be equal to h ensure sufficient history is available for the first prediction.
To thoroughly evaluate the effectiveness and robustness of PDT under different training configurations. we conduct comprehensive experiments across different learning rates from 0.001 to 0.1 (0.001, 0.01, 0.05, 0.1) and batch sizes from 32 512 (32, 64, 128, 256, 512). All experiments were repeated with five random seeds (0, 100, 200, 300, 400) to ensure statistical significance. All experiments are performed on AeexNet with the CIFAR-I0 dataset, using SGD as the baseline optimizer and trained for 60 epochs. The PDT-related hyperparameters mentioned in Sec. were prediction interval=1, 1, start epoch=5, and past snapshot counts=5.
were set to prediction step=5,
2.10
2.
2.5
1.8
1.2
Pedines:
Reading-2
Prodo
1.5
3.4
Frade
Budi
500
Epach
(a) 1r=001, in start=10, 10
(b) Ir=0.05, step=9,
(c) 1r=00001, int=1, step=5, snap=5.
(d) Ir=0.01, step=S, int=1, start=10..
start=10, 10, snap=10.
Figure 11:
The influence of different parameters. (a) prediction steps, (b) prediction interval, (c)
starting epoch, (d) past snapshot counts. Trained on CIFAR-10 using AlexNet, batch size=256.
The results in Table Ashow the impact of different batch sizes and learning rates on the performance of PDT. At lower learning rates (0.001, 0.01, and 0.05), PDT consistently outperforms SGD in terms of convergence speed across different batch sizes. PDT shows a significant reduction in the runtime to reach baseline best loss, with an average runtime reduction of 72.76% compared to SGD. For higher learning rates (0.1), I), both SGD and PDT struggled to achieve stable training, and PDT's advantage over SGD became less pronounced. Sometimes PDT can significantly reduce the convergence time (for example, when batch size 64), but other times the accuracy will drop significantly after reaching result in gradient explosion. This suggests that the high learning rate introduced significant stochasticity, reducing the effectiveness of PDT's prediction mechanism. Smaller batch sizes (32, 64) generally achieve more significant runtime reductions.
high point,
To address the stability issues observed at higher learning rates and larger batch sizes, different from the previous fixed learning rate, we investigated the effectiveness of the learning rate scheduler. We tested the Cosine Annealing learning rate scheduler with a minimum learning rate of le-3. Taking batch size 256 as example, we observe significantly improved stability and performance. The results are shown in Table 5) The results are particularly noteworthy at higher learning rates (Ir=0.1), 1), where the previous experiments in Table 4 show considerable variance. With the cosine annealing scheduler, PDT achieves consistent accuracy improvements across all learning rates while maintaining substantial runtime reductions.
To further investigate PDT's compatibility with different optimization methods, we compare its performance when integrated with different optimizers (SGD, SGD with momentum, and Adam) while keeping the network architecture and other configurations fixed. For SGD with momentum, we set the momentum factor to 0.9. All experiments are conducted on AlexNet with CIFAR-10 100 using batch size 256, maintaining the same PDT hyperparameters as in previous experiments. The learning rate is 0.1 for SGD, 0.001 for SGD with Momentum, 0.0005 for Adam. The results are shown in Table
20

---

Published as a conference paper at ICLR 2026
Table 4: Impact of learning rates and batch sizes on PDT performance. Trained on CIFAR-10 using AlexNet. Note: bold numbers indicate the best performance and underlined numbers indicate the second best performance for each column.
Batch Size
Final Accuracy (mean std)
Best Loss
Time to Baseline Best Loss (s) (mean : std)
Runtime Reduction (%)
1r
Method
(mean 1 std)
SGD
0.6981 + 0.0458
0.6376 0.0127 0.2724 + 0.0166
1232.29 4.45
0.001
40.64
PDT
0.6903 + 0.0885
731.52 + 12.84
SGD
0.8118 + 0.0041
0.0046 + 0.0008
1194.89 + 21.09
0.01
24.25
32
PDT
0.8146 1 0.0048
0.0021 + 0.0012
905.07 1: 120.51
SGD
0.8049 +: 0.0053
0.0156 0.0029
1180.72 12.31
0.05
64.57
PDT
0.8020 +: 0.0052
0.0149 + 0.0073
418.38 + 0.00
SGD
0.1000 + 0.0000
0.3346 0.0098
1172.49 39.08
0.1
PDT
0.1000 + 0.0000
0.3364 + 0.0132
SGD
0.5384 + 0.0173
1,2295 + 0.0261
902.16 + 19.68
0.001
35.82
PDT
0.5329 1 0.1152
0,8798 + 0.0257
578.99 + 55.74
SGD
0.7850 1 0.0226
0.0087 + 0.0030
800.35 5.39
0.01
23.32
64
PDT
0.8140 + 0.0021
0.0015 + 0.0010
613.70 + 8.80
SGD
0.8067 + 0.0035
0.0051 0.0016
798.20 3.50
0.05
27.54
PDT
0.8029 + 0.0029
0.0045 + 0.0006
578.36 + 16.48
SGD
0.6442 + 0.2733
0.0484 + 0.0522
910.37 = 18.03
0.1
56.23
PDT
0.7976 + 0.0033
0,0218 + 0.0011
398.48 + 21.34
SGD
0.2882 + 0.0212
1.8456 + 0.0300
812.42 + 21.20
0.001
17.48
PDT
0.2951 + 0.0440
1.6972 + 0.0272
670.37 + 23.93
SGD
0.7825 + 0.0065
0.0675 + 0.0052
661.09 + 6.35
0.01
14.68
128
PDT
0.8009 + 0.0062
0.0058 + 0.0008
564.02 + 16.35
SGD
0.7969 + 0.0093
0.0039 + 0.0017
662.48 + 7.73
0.05
9.15
PDT
0.8011 1: 0.0067
0,0016 + 0.0017
601.86 1: 17.78
SGD
0.7916 1 0.0027
0.0083 1 0.0014
803.93 3.07
0.1
8.20
PDT
0.7863 1: 0.0087
0.0096 1 0.0016
737.97 1 0.00
SGD
0.1171 + 0.0092
2.2991 + 0.0011
747.83 + 20.30
0.001
7.08
PDT
0.1453 + 0.0213
2.2979 = 0.0026
694.91 = 14.63
SGD
0.6989 + 0.0301
0.5814 + 0.0147
660.37 + 0.71
0.01
19.98
256
PDT
0.7450 + 0.0236
0.1855 0.0172
528.41 7.26
SGD
0.7931 + 0.0034
0,0004 + 0.0003
648.39 + 8.57
0.05
21.71
PDT
0.7916 + 0.0016
0.0015 + 0.0014
507.62 + 11.36
SGD
0.3742 + 0.3359
0.0508 + 0.0576
771.77 + 3.06
0.1
PDT
0.3796 + 0.3425
0.0012 + 0.0011
SGD
0.1170 1 0.0251
2.3017 1 0.0005
748.44 1. 42.46
0.001
6.23
PDT
0.1377 1 0.0288
2.3020 0.0001
701.82 1: 23.31
SGD
0.5710 +: 0.0203
1.1920 +. 0.0238
671.28 1 9.03
0.01
18.89
512
PDT
0.5985 1 0.0078
0.0252 0.8311
544.46 +: 12.10
SGD
0.7717 + 0.0038
0.0311 + 0.0174
668.59 + 7.30
0.05
10.11
PDT
0.7669 + 0.0237
0.0034 0.0014
601.01 + 44.11
SGD
0.3721 + 0.3332
0.0648 + 0.0735
768.97 + 3.12
0.1
PDT
0.4420 + 0.3420
0.0373 0.0155
21

---

Published as a conference paper at ICLR 2026
Table 5: Impact of learning rates PDT performance. Trained CIFAR-10 using AlexNet, batch size=256, with CosineAnnealingLLR scheduler, minimum learning rate 1e-3. Note: bold numbers indicate the best performance and underlined numbers indicate the second best performance for each column.
Batch Size
Final Accuracy (mean std)
Best Loss
Time to Baseline Best Loss (s) (mean = std)
Runtime Reduction (%)
1r
Method
(mean + std)
SGD
0.1217 + 0.0126
2.2991 0.0011 2.2980 + 0.0025
757.66 + 26.54 682.79 2.13
0.001
9.88
PDT
0.1461 1 0.0213
SGD
0.6451 +. 0.0102
0.9276
0.0212
745.97 1 47.19
0.01
41.54
256
PDT
0.6974 1 0.0073
0.5853 i 0.0159
436.07 + 16.09
SGD
0.7852 + 0.0016
0.0020
0.0001
675.04 + 27.56
0.05
37.13
PDT
0.7936 + 0.0030
0.0006 + 0.0001
424.39 + 20.40
SGD
0.7930 + 0.0023
0.0002
0.0000
665.27 9.08
0.1
19.67
PDT
0.7978 + 0.0032
0.0002 + 0.0000
534.41 + 12.64
Table 6: Impact of baseline optimizers (SGD, SGD with Momentum, and Adam) PDT performance. Trained on CIFAR-10 using AlexNet, batch size=256, momentum=0.9, with Annaaln scheduler. Note: bold numbers indicate the best performance and underlined numbers indicate the second best performance for each column.
Final Accuracy (mean std)
Best Train Loss
Time Baseline Best
Runtime Reduction (%)
Ir
Method
(mean + std)
Loss (s) (mean std)
SGD PDT
0.7930 0.0023 0.7978 + 0.0032
0.0002 0.0000 0.0002 + 0.0000
665.27 9.08
0.1
19.67
534.41 + 12.64
Momentum
0.6672 + 0.0068
0.8609
0.0166
752.74 9.62
0.001
41.06
PDT
0.7298 + 0.0051
0.5358 + 0.0165
443.68 8.75
Adam
0.7952 + 0.0063
0.0001
0.0000
779.13 + 11.81
0.0005
14.87
PDT
0.8050 + 0.0050
0.0002 + 0.0000
663.28 + 15.30
22

---

Published as a conference paper at ICLR 2026
A.7 ANALYSIS OF MASK DISTRIBUTION
We further analyze the mask distribution and dynamics in AlexNet. Fig. 12 shows how the ratio of the predicted weights evolves over training epochs. The analysis is conducted using the same experimental setup as in Fig. b)),), where AlexNet is trained on CIFAR-10.
CrmO
CEV uves
Gond
1pach
(a)
(b)
Figure 12: Analysis of mask distribution in AlexNet. (a) Layer-wise mask evolution over training
epochs. (b) Comparison of prediction ratios between convolutional and fully connected layers.
Overall: Predict=18.14 SGD=81.84% SGD Update
Peddct Update
Layer name: features Bueright (Predict-31. 30.3 SGD- 62%))
10000
15000
20000
Weegtt Indes
Layer name: features. weight (Phedict25..d 25.3%% 54 82%))
100000
200000
3000C0
400000
500000
GCO000
features (Phedict=22, SGDm 12.0m12.788. 7.78%)
100000
200000
300000
400000
500000
600000
70000D
8C0000
Vie git Indix
Layer name: featurres 10. weight (Predict=27 57%, 5GD-122.333 72.33%
100000
200000
500980 Weight Incee
Layer name: classiher: swwight rediccn 48%, .4444
Layyer
0.4
0.6
0.8
1.2
1.4
1.6
Veight
167
Layer classii en wdde:w27. 30%, SGD0 72.20%
5000
10000
15000
20000
25000
30000
35200
10000
Vie gtt Incex
Figure 13: A snapshot at epoch 20 with the mask heatmap for different layers.
Fig. 12 a) presents the layer-wise evolution of the ratio of the predicted weights throughout the training process. We observe pattern here: the masked ratio of each layer starts relatively high, maintaining stable period, and then gradually declining. The decline phase at the later epochs suggests that as the network approaches convergence, it relies more on gradient- based updates rather than predictions. This aligns with the intuition that predictive updates be beneficial in the early
23

---

Published as a conference paper at ICLR 2026
phases for accelerating convergence but become less necessary the model stabilizes. The early convolutional layers (e.g., ConvO) exhibit more fluctuations in the percentage of predictive updates, suggesting higher sensitivity to training dynamics.
Fig. 12 b) tracks the evolution of predicted weights ratios by layer type. The overall percentage of predictively updated weights is also included. Interestingly, convolutional layers consistently maintain a higher prediction ratio compared to fully connected layers throughout the training process. Due to the majority of the weights in the AlexNet network belonging the fully connected layers (54.6 million vs. 2.5 million), the overall masked ratio closely follows the trend of fully connected layers.
To provide a finer- sner-grained visualization of the mask distribution, Fig. 13 depicts the mask heatmap for different layers at epoch 20. Each horizontal band represents a layer, where blue regions indicate weights updated by SGD, and red regions correspond to weights updated by prediction results. We can observe that the distribution of predictive updates is not uniform across layers, with some layers showing clustered regions of predictive updates, potentially indicating structured weight adaptations.
A.8 PROFILING AND OVERIEAD ANALYSIS
To validate our theoretical complexity analysis and quantitatively measure the overhead, we conduct comprehensive profiling experiments across three diverse architectures: AlexNet (CNN, 57M pa- rameters) on CIFAR- 10, ResNet-50 (deep CNN, 25.6M parameters) on ImageNet- Ik, and ViT-Base (Transformer, 86.4M parameters) on ImageNet-1k, We measure wall-clock time, GPU memory usage, and FLOPs for both the baseline optimizer and PDT.
A.8.1 MEASUREMENT METHODOLOGY
Memory Overhead. The Peak GPU memory captures true peak GPU usage during training, measured via terch. cuda... max memory al llocated (). Currently Allocated is the memory actively allocated at the end of each epoch. Reserved by Allocator is the total memory reserved by PyTorch's caching allocator. The PDT Memory Overhead consists of Snapshot Storage (h copies of weights), SVD Workspace (the temporary storage for intermediate matrices), and Other (includes optimizer momentum buffers, weight cloning buffers, temporary tensors, and memory fragmentation). The PDT memory overhead ratio is computed the total PDT overhead divided by the baseline's peak GPU memory.
Runtime Overhead. The Total PDT Runtime Overhead is defined as the difference between the PDT epoch time and the baseline SGD epoch time. We further decompose this into Core PDT Operations (including SVD decomposition, DMDD-based weight prediction, and gradient masking) and Auxiliary Operations (including weight cloning for snapshot storage, CUDA synchronization. and data movement between CPU/GPU). The PDT runtime overhead ratio computed as total PDT overhead divided by total epoch time.
FLOPs Overhead. Total floating boating-ppint operations for the entire training process.
All profiling experiments are conducted on a single GPU (Nvidia RTX A6000 or H100) to ensure accurate measurement of peak memory and timing without the interference of distributed commu- nication overheads. For PDT, we use past snapshot counts of h = 5 and a prediction interval of T;
A.8.2 PROFILING RESULTS
1. AlexNet on CIFAR-10. Table summarizes the results for AlexNet on CIFAR-10. 10. The peak GPU memory usage for PDT increases by approximately 8.4 GB compared to the baseline. This absolute increase is primarily due to the snapshot storage (1.1 GB for h = 5) and the SVD workspace (1.7 GB for intermediate matrices). The remaining overhead (~ 5.6 GB) is attributed to temporary buffers and memory fragmentation during the weight update process. While the relative increase (1 (145%) appears large due to the small baseline footprint of AlexNet on CIFAR- 10, the absolute peak usage (9.1 GB) fits comfortably within the capacity of modern GPUs. This confirms that the O(N x h) space complexity is practical, as described in Appendix The total PDT overhead is approximately 333 ms per epoch, with the PDT overhead ratio for only 4.30% % Decomposing this overhead reveals
24

---

Published as a conference paper at ICLR 2026
that core PDT operations account for 194 (58.3%), while auxiliary operations (weight cloning and synchronization) account for 138.8 ms (41.7%). Notably, the SVD operation itself takes only theoretical claim that the SVD cost is computationally efficient
98 ms, empirically confirming even for frequent predictions. The additional FLOPs introduced by PDT are negligible. The total training FLOOPs increase from 3.4043 x 10t5 (baseline) 3.4047 x 10 (PDT), 0.012% increase, validating that the computational cost is dominated by gradient calculations in the baseline optimizer rather than in the PDT operations.
Table 7: Memory and Runtime overhead of PDT compared to baseline. Trained on CIFAR-10 using AlexNet, batch size=256, SGD baseline optimizer with Ir-0.05, with CosineAnnealingLR scheduler. All experiments ran on single Nvidia RTX A6000 (48 GB) GPU.
Metric
Baseline
PDT
GPU Memory Usage
Peak GPU Memory (MB)
735.9
9158.2
Currently Allocated (MB)
468.7
5078.0
Reserved by Allocator (MB)
994.0
11454.0
PDT Memory Overhead
Total PDT Overhead (MB)
8422.4
Overhead Ratio (%)
1145%
PDT Memory Overhead Breakdown
Snapshot Storage (MB)
1088.0
SVD Workspace (MB)
1740.9
Other (MB)
5593.4
PDT Runtime Overhead (per epoch)
Avg Time per Epoch (s)
7.465
7.746
Total PDT Overhead (s)
0.333
_
PDT Overhead Ratio (%)
4.30 %
PDT Runtime Overhead Breakdown
Core PDT Operations (ms)
194.0
SVD decomposition (ms)
98.3
Prediction (ms)
88.1
4
Masking (ms)
7.6
Auxiliary Operations (ms)
138.8
2. ResNet-50 on ImageNet-1k. To further validate overhead analysis, we conducted comprehensive profiling on ResNet-50 + ImageNet-Ile Ik with varying batch sizes (64, 128, 256). Table S presents detailed memory and runtime measurements for both baseline and PDT. The absolute total PDT memory overhead remains nearly constant (~ 1.46 GB) across all batch sizes, validating that PDT's space complexity ONN x h) is independent of data scale. Consequently, the PDT memory overhead ratio decreases monotonically as baseline memory increases with batch size from 47.9% (bs=64) to 13.0% (bs=256). In order to run on single GPU. For ResNet-50 and ViT results reported in Table ! we used a batch size of 1800 (600 per GPU, trained on GPUs), so we estimate the overhead ratio should be even lower. In some cases in Table g the Avg. Epoch time of PDT was slightly shorter than baseline, likely due to measurement variance. The total PDT runtime overhead (90-1 110 ms) consistent and negligible. The core PDT operations (SVD + prediction + masking) take only ~ 51 ms per epoch, which is insignificant compared to the (635s-804s) epoch time.
overhead more accurately, all experiments in this section were
3. ViT-Base on ImageNet-lk. To validate the generality of our overhead analysis across different architectures, we conducted profiling on ViT-Base (Transformer architecture) with ImageNet- at batch size of 256. Table T presents detailed memory and runtime measurements for both baseline and PDT.
Despite ViT-Base having significantly more parameters (86.4M) than ResNet-50 (25.6M), the memory overhead breakdown remains consistent with our theoretical analysis: snapshot storage and SVD workspace dominate (40% and 53% respectively), with only 7% attributed to auxiliary memory. The absolute memory overhead (4,954 MB) is larger than ResNet-50 (1,463 MB) as expected from the O(N h) scaling. The runtime overhead remains negligible at 0. 11%.
25

---

Published as a conference paper at ICLR 2026
Table Memory and Runtime overhead PDT compared baseline across different batch sizes. Trained on ImageNet-Il Ik using ResNet-50, SGD with Momentum baseline optimizer with Ir=0.1. 1, momentum=0.9, with CosineAnnealingL.. scheduler. All experiments ran on single Nvidia H100 (80 GB) GPU.
Baseline
PDT
Metric
bs=64 hs= bs=256
hs=64
bx= 128
bs=256
GPU Memory Usage
Peak GPU Memory (MB)
3,062
5,792
11,258
4,528
7,255
12,721
Currently Allocated (MB)
400
436
512
2,439
2,436
2,439
Reserved by Allocator (MB)
3,404
6,766
13,424
7,392
10,864
15,298
PDT Memory Overhead
Total PDT Overhead (MB)
_
1,466
1,463
1,463
Overhead Ratio (%)
47.9%
25.3%
13.0%
PDT Memory Overhead Breakdown
Snapshot Storage (MB)
585
585
585
SVD Workspace (MB)
780
780
Other (MB)
101
98
98
PDT Runtime Overhead (per epoch)
Avg Time per Epoch (s)
813.9
653.1
626.5
804.0
651.5
635.3
Total PDT Runtime Overhead (ms)
90
90
110
_
Overhead Ratio (%)
0.011 %
0.013% %
0.017% %
PDT Runtime Overhead Breakdown
Core PDT Operations (ms)
51.5
50.7
51.1
SVD decomposition (ms)
31.0
30.7
4
Prediction (ms)
19.7
19.5
19.6
4 Masking (ms)
0.8
0.8
0.8
Auxiliary Operations (ms)
38.5
39.3
58.9
A.8.3 SCALABILITY AND GENERALIZATION ANALYSIS
The profiling results align robustly with theoretical analysis in Appendix A.3 confirming that PDT is scalable to large modern architectures.
AlexNet has large number of parameters (57M) relative the small size of the CIFAR-10 dataset. This results in scenario with high memory requirements for SVD (proportional to N) but very short epoch duration (low compute load). For large-scale dataset like ImageNet, the computational load epoch increases dramatically, while the parameter count N (which dictates PDT overhead) does not change with a specific network architecture. The disparity of the memory overhead ratio (AlexNet: 1.145%, ResNet-50: 13%, ViT-Base: 26%) can be explained by the composition of baseline memory. For AlexNet on CIFAR-10 10, the baseline memory (736 MB) is dominated by model parameters because the input image size (32 x 32) generates very small activation maps. For ImageNet-IK 1k with (224 x 224) inputs, the baseline memory (1 1-91 GB) is dominated by activation memory, while parameter storage is smaller fraction. Since PDT overhead scales only with parameters and not activations, the relative overhead drops dramatically.
In conclusion, while PDT introduces linear space complexity O(N x h), this overhead is amortized a large-scale training scenario where memory dominated by activations and runtime is dominated by gradient computation.
26

---

Published as a conference paper at ICLR 2026
Table 9: Memory and Runtime overhead PDT compared to baseline. Trained on ImageNet- Ik using ViT-Base, batch size=256, AdamW W baseline optimizer with lr-0.0066, weight decay=0.05, with Cosine scheduler. All experiments ran on a single Nvidia H100 (80 GB) GPU.
Metric
Baseline
PDT
GPU Memory Usage
Peak GPU Memory (MB)
19,022
23,976
Currently Allocated (MB)
1,536
8,406
Reserved by Allocator (MB)
19,462
35,328
PDT Memory Overhead
Total PDT Overhead (MB)
4,954
Overhead Ratio (%)
26.0%
PDT Memory Overhead Breakdown
Snapshot Storage (MB)
1,981
SVD Workspace (MB)
2,642
Other (MB)
330
PDT Runtime Overhead (per epoch)
Avg Time per Eppoch (8)
831.3
832,1
Total PDT Runtime Overhead (ms)
877
Overhead Ratio (%)
0.11% %
PDT Rumtime Overhead Breakdown
Core PDT Operations (ms)
193.2
SVD decomposition (ms)
127.4
Prediction (ms)
62.8
Masking (ms)
3.0
Auxiliary Operations (ms)
683.6
Table 10: Comprehensive masking ablation results on AlexNet trained CIFAR-IO. Success rates are shown as successful/total runs. Final accuracy reported as mean + std across five seeds (successful runs only).
Contiguration
LR=0.01
LR=0.05
Overall
Success Rate
Final Acc (%)
Success Rate
Final Ace (%)
Success
Baseline (SGD)
64.40 0.71
5/5
78.68 0.21
10/10
5/5
58.04 + 1.29
5/5
78.23
0.64
Accel Only
16/20
48.20 + 8.19
2/5
10.00 0.00
3
0/5
Crashed Crashed
1/5
10.00 Crashed
Consistency Only
1/20
5
0/5
0/5
69.39
0.79
5/5
78.98
0.18
Full PDT
20/20
5
5/5
69.10 + 1.33
5/5
79.48 + 0.20
A.9 ABLATION STUDY: MASKING CRITERION ANALYSIS
we introduce a masking mechanism including two principles: the acceleration effective- ness criterion (Eq. 6) and the dynamic consistency criterion (Eq. 7 To investigate the contribution and distinct role of each masking criterion, we conducted a comprehensive ablation study AlexNet trained on CIFAR-10, 10, using SGD as the base optimizer. We compare four configurations: (1) Baseline (standard SGD), (2) Accel Only (apply only acceleration effecttivnness criterion for masking), (3) Consistency Only (apply only dynamic consistency criterion for masking), and (4) Full PDT (our complete PDT method with both criteria). Each configuration evaluated with learning rates 0.01 and 0.05, prediction steps T E {3, 5). and five random seeds (0, 100, 200, 300, 400), totaling 70 experimental runs.
In Sec. 3.2
Table 10 presents the success rate (proportion of runs that completed without crashing) and final validation accuracy (computed only from successful runs) for each configuration. Consistency Only
27

---

Published as a conference paper at ICLR 2026
achieves only successful run (seed=300, LR=0.05, 3) with 10% aceuracy (random level), demonstrating catastrophic failure. Accel Only achieves 80% success but shows particular instability at 5 with LR=0.05 (2/5 success), with successful runs achieving only 10% accuracy, indicating its inability to predict longer steps. The full PDT achieves perfect robustness (20/20 success) and the highest validation accuracy both learning rates, outperforming the baseline performance, and validating the necessity of combining both criteria.
Figure 14 visualizes the training dynamics across different learning rates and prediction steps (7). Since there is no need to adhere to the dynamic consistency criterion, Accel only consistently has higher masked ratio than the full PDT. In all configurations, its train loss is higher than that of full PDT. Specifically, in Fig. [4fd) (LR=0.05,T = 5), the Accel Only validation accuracy collapses to random chance (10%) after the initial epochs. Without the consistency criterion, too many predictions with incorrect directions are accepted, which causes gradient explosion, especially when predicting multiple steps. In contrast, Full PDT consistently achieves the lowest training loss and highest validation accuracy across all configurations. The shaded variance regions for Full PDT are notably tighter than those for Accel Only, demonstrating that combining both criteria makes the training process significantly more stable.
2.
AccE Only (UNd]
Kuualise Crcs load?
Aecaldes Full POT (wo))
Stalled Racco
-0
0.5
10
20
40
37
55 Epech
43
10
15
Epach
Spoch
(a) LR=0.01, prediction step (r)
Bnseline Stet
AccetCes (e-5)
Wakcation Acuracy (044
posted Rasto natto M6))
Thirkng
03
50
20
(b) LR=0.01, prediction step (7) 55
Exseiine tre (500) full rowy
Accal On Toad
Masied Tallo (56
Aecerde (1-1) FOT (1H2)
Epoch
Epech
Epoch
(c) LR=0.05, prediction step (r) =
Aecatoes Full POT (1w))
ROCE Onyy (Uwd)
3.5
Deseline (SG0)
Coalted Para
Tharteg
ot (uch)
3.0
-5
25
40
30
15
53
Epoch
Epoch
Spoch
(d) LR=0.05, prediction step (7) 5.
Figure 14: Mean training curves (with standard deviation indicated by Shaded regions) and perfor- mance comparison of bascline and pdt with different masking criteria. Trained on CIFAR- 10 using AlexNet, batch size 256, with CosineAnnealingLLR scheduler. Only successful runs are included.
To further investigate the contribution of each masking criterion across different training stages, we decompose the predictions rejected by PDT and examine which criterion rejects them. Figure T5
28

---

Published as a conference paper at ICLR 2026
shows the temporal evolution of each criterion's contribution throughout training. We analyze the Full PDT configuration (LR=0.05) across five random seeds to understand how the acceptance and rejection ratios evolve. Experimental results from ==3 and r=5 show similar pattem. Ratio of predictions rejected by Consistency Criterion (Eq. 7 due to opposite direction remains stable at ~50 throughout training. Ratio of predictions rejected by Acceleration Criterion Lower Bound (Eq. (for being too small (< 1x SGD update)) decreases to near zero, indicating this bound is primarily active early in training. Ratio of predictions rejected by Accel Upper Bound (for being too large (> (T)x SGD update)) increases dramatically (from 20% to ~50%) becoming critical in the later training stage.
The overall mask ratio decline is due to the increasing rejection by the upper bound of the acceleration effectiveness criterion. The gradual decline in mask ratio reflects the changing optimization land- scape: early in training, when the loss landscape is steep and gradients are large, DMD) sometimes underestimates the required step size. However, as training progresses and gradients shrink near convergence, DMD predictions rarely fall below the minimum threshold. When near convergence, gradients become small, noisy, and oscillatory. The upper bound becomes critical for preventing divergence by rejecting these over -gggessgee predictions.
M111 Aucep Aodita (olaeed aegecned Accel teilr Lower loud Rupotiin Aud Upe
M111 cegilnilikai (RATicoinin betiaaa segected Acdel (erie Lowe' Jond Pupiticn Aud Upe
10
20
50
10
30
50 60
(a) LR=0.05, prediction step 3r)=3. (b) LR=0.05, prediction step (7) 55
Figure 15: Temporal evolution of masking criteria contributions during training. Results from full PDT (both criteria applied). Trained on CIFAR- 10 using AlexNet, LR=0.05, batch size 256, with CosincAnnealingL.R scheduler. Only successful runs are included.
From the above ablation results, compared to a single criterion, the complete PDT with both criteria has smooth loss curves, faster convergence, and low variance across seeds. This provides strong empirical evidence supporting our theoretical design and demonstrates that neither criterion alone sufficient. Both criteria play distinct, complementary roles. The dynamic consistency criterion provides stable, stage-agnostic directional filtering. The acceleration effectiveness criterion ensures predictions provide speedup by enforcing magnitude bounds. During the early training stage, the lower bound guarantees that each accepted prediction moves the weights at least as far as a single SGD step, enabling acceleration. The upper bound prevents over-aggressive updates that could cause divergence and ensures stability throughout the training process.
A.10 IMPACT OF DIFFERENT INITIAL LEARNING RATES ON MASKED RATIO
To investigate the impact of different initial learning rates on masked ratio behavior and computational efficiency, we conducted systematic experiments on AlexNet trained on CIFAR-IO. We compare Baseline (SGD) and PDT (prediction steps r = 5, prediction interval T; = 1, past snapshot counts
5. starting epoch To 5) across four different learning rates: 0.001, 0.01, 0.05, and 0.1. All experiments use batch size 256, CosineAnnealingl!
scheduler with Irmin = 10-3 and train for 60 epochs. Each configuration was evaluated with three random seeds (0, 100, 200). We measure mask acceptance ratio by training stage (carly: epochs 5-20, mid: epochs 21-40, late: epochs 41-60) and final validation accuracy.
As shown in Table 11 larger learning rates lead to higher mask acceptance ratios during early training. Specifically, LR=0.1 achieves the highest early-training masked ratio (18.77%), followed by LR=0.05 (16.94%), LR=0.01 (8. 18%), and LR=0.001 (0.06%). For relatively large learning rates (e.g., 0.05 and 0.1), this trend is not strictly observed. As training progresses, the masked ratios across different learning rates decline to varying degrees. However, for extremely small learning rates (e.g., 0.001),
29

---

Published as a conference paper at ICLR 2026
the masked ratio remains consistently low. PDT consistently improves over baseline across all learning rates, demonstrating robustness.
Table 11: Mask acceptance ratio and validation accuracy across learning rates: mask acceptance ratio by training stage and final test accuracy (mean + std across 3 seeds). Training stages: Early (epochs 5-20), Mid (epochs 21-40), Late (epochs 41-60).
LR
Early (%)
Mid (%)
Late (%)
Overall (%)
Base Acc (%)
PDT Ace (%)
0.001
0.06 + 0.00
0.07 = 0.00
0.10 + 0.02
0.08 + 0.01
11.83 + 1.04
14.92 + 2.44
0.01
8.12 + 0.20
7.10 = 0.06
2.77 + 0.02
5.77 + 0.07
63.96 + 0.14
68.43 + 0.64
0.05
16.94 0.20
8.52
0.68 + 0.15
7.94 + 0.31
78.67 0.13
79.24 + 0.22
0.1
18.77 = 0.24
7.10 = 0.52
0.37 + 0.05
7.91 = 0.20
79.76 + 0.27
80.06 0.20
LR-0.001
LR-0.01
Lt-i0.0x
Masied Raboo (%)
10-
Epoch
Figure 16: Mask acceptance ratio evolution across different learning rates (mean + std shading across 3 seeds). Trained on CIFAR-10 10 using AlexNet, batch size 256, with CosineAnnealinggl scheduler.
Figure 16 visualizes the mask ratio evolution across different learning rates. Early training exhibits the highest acceptance, as the loss landscape is steep and gradients provide a strong signal. By late training (epochs 41-60), mask ratios converge to low values across all learning rates.
LR=0.001
LR=0.01
100
Accepted (Both Satisfied)
Accepted (Soth tenia Satisfied)
Rejected: Criterion
Rejectedi Criteron
Rejected: Upper Bound
Rejocted: Upper Bound
Radio (35)
Rotio (36)
40
20
20
Epoch
LR=0.05
LRt0.11
100
100
PDT Accepted (Both Cnter Satistied;
PDT Accepted (Botn Cr tena Satistid)
Rojected: Consictency Critarin ion
Rejoctedi Consistoncy Critrion
Rejected: Accel Criterion Lower Bound
00
Rejected: Accel Crtteiion Lower Bound
Crieria
Rad0 150)
Ratoo 196)
20
20
10
30
Esoch
Figure 17: Temporal evolution masking criteria contribution across learning rates. Trained on CIFAR-10 using AlexNet, batch size 256, with Cosine scheduler.
30

---

Published as a conference paper at ICLR 2026
The temporal evolution of masking criteria contribution shown in Fig. 17 reveals why extremely small learning rates (e.g., 0.001) lead to consistently low mask acceptance ratios during the entire training stage. The mask acceptance behavior is governed by the interaction between the learning rate and the Acceleration Effectiveness Criterion (Eq. 6 At LR=00001, the figure shows that predictions are overwhelmingly rejected by the consistency criterion and the Upper Bound of acceleration criterion (Green line). When the learning rate is extremely low (e.g., 0.001), the gradient steps are small, making the "allowable acceleration window" microscopic. Although DMD predicts a trajectory based on historical dynamics, the magnitude of this prediction, even if small, easily exceeds the excessively strict upper bound imposed by the tiny learning rate. contrast, at LR=0.05 and 0.1, the upper bound constraint is relaxed. The gradient signal is strong enough that the SGD step size is comparable to the DMD prediction scale. Consequently, larger proportion of predictions fall comfortably between the lower bound and the upper bound, resulting in significantly higher masked ratio.
In conclusion, the stage-wise analysis reveals that masked ratios are highest during early training and decline adaptively toward convergence, prioritizing stability over acceleration in the final stages. Contrary to the intuition that smaller steps imply stability and higher acceptance, our analysis confirms that larger learning rates necessary to generate the clear dynamic signals required for high-iidlityy DMD predictions. PDT is most effective when the baseline optimizer takes steps large enough to define "permissible region" that accommodates the scale of DMD's spectral predictions.
A.11 EFFECT OF NON-L.LD. TRAINING DATA
We further investigate the robustness of PDT under some challenging training conditions. For example, when the batch is too small for a diverse dataset like ImageNet, the weight updates could be chaotic since each consecutive batch is no longer an identical distribution. There are two experimental designs that can test this: 1) test PDT on a very large dataset like ImageNet-224 and 2) design batching scheme to intentionally violate the i.idd assumption of mini-batches using a smaller dataset such CIFAR- 10. In the second design, we maintain the normal batch size, but only put samples of the same class the batch. We also randomize the batch sequence instead of using any fixed order so that there is no regular training set dynamics that DMD might pick up on.
5G0
0.8
2.0
0.7
3.5
1.0
0.6
0.15
3.0
0.5
7ain
0.4
0.3
0.09
0.5
Eaech
(a) Under normal i.i.d. training data distribution.
5G0
0.150
0.125
0.100
Medked
2.0
Tet
0.3
0,025
0.3
0.000
Epoch
Epach
(b) Under non-i.i.d. training data distribution.
Figure 18: Performance comparison between SGD and PDT under i.i.d. and non-i.i.d. training data distributions, with the same hyperparameters configuration. Trained on CIFAR- 10 using AlexNet, batch size=128, Ir=0.05, with CosineAnnealinggl, scheduler. The shaded areas represent the standard deviation across 5 runs with different random seeds (0, 100, 200, 300, 400).
Fig. 18 and Table 12 show the performance and runtime comparison between SGD and PDT under the non-i.i.d. setting using the second experimental design since non-i.i.d. is guaranteed. We preserve
31

---

Published as a conference paper at ICLR 2026
Table 12: Performance and runtime comparison between SGD and PDT under i.i.d. and non-i.i.d. training data distributions, with the same hyperparameter configuration. Trained on CIFAR-10 using AlexNet, batch size=128, 1r=0.05, with CosineAmeellingL. scheduler.
Training Data Distribution
Final Accuracy
Best Train Loss
Time to Baseline Best Loss (s) (mean std)
Runtime Reduction (%)
Method
(mean std)
(mean + std)
SGD
0.7969 + 0.0093
0.0039 0.0017
662.48 7.73
i.i.d.
9.15
PDT
0.8011 + 0.0067
0.0016 0.0017
601.86 = 17.78
SGD
0.7067 + 0.0062
0.1053
0.0874
806.83 13.15
non-ii.d.
27.90
PDT
0.7159 1 0.0103
0.0119 0.0057
581.73 19.34
the original i..d. sampling of the validation set. All experiments are repeated with five random seeds (0, 100, 200, 300, 400) to ensure statistical significance.
We make some interesting observations. First, despite the challenging non-i.id. setup, PDT still achieves better performance than SGD in terms of faster convergence without sacrificing accuracy. However, we also observe that in the non-i.i.d. case, learning starts out much more slowly for both SGD and PDT and both take longer to converge. Second, in the non-i.i.d. case, the variance of each of the performance curves is generally larger than those of the i..d. case. This is because the model needs to handle more abrupt transitions between different class distributions.
Fig. 18 and Table 12 further demonstrate that PDT's advantage extends beyond standard i.i.d. training conditions, showing its robustness to challenging data sets where traditional assumptions about data distribution are violated.
A.12 CROSS-DOMAIN EVALUATION ON NATURAL LANGUAGE PROCESSING
To demonstrate the generalization of PDT beyond computer vision, we evaluate its effectiveness on text classification tasks. This experiment validates PDT's applicability fundamentally different data modality (discrete text) and architecture (Recurrent Neural Networks).
We employ a deep LSTM network for 4-class topic classification on the AG News dataset (Zang et al. 2013; Gulli). The architecture consists of an embedding layer, a 4-layer stacked LSTM (512 hidden units per layer), and linear classifier, totaling 8.25M parameters. This setup differs significantly from CNNs, particularly in terms of gradient flow dynamics (backpropagation through time). We use SGD with learning rate of 0.1 and batch size 128 for 30 epochs. The PDT hyperparameters identical to those used in vision experiments (T 5, h 5, To 5, T 1) to assess the robustness of PDT without domain -specffc tuning. Results are averaged over 3 random seeds.
Table 13: Performance comparison on AG News dataset (text classification) using a deep LSTM model. Trained single Nvidia RTX A6000 GPU for 30 epochs, 1r=0.1, batch size=128, with CosineAnnealingLRR scheduler.
Runtime Reduction (%)
Optimizer
Final Accuracy
TTB- (s)
TTB-Ace (s)
Train Loss
Val. Acc.
SGD PDT
86.10 + 1.30
700.8 + 72.1
694.4 + 76.9
26.1
20.5
88.44 0.59
517.9 + 102.2
551.7 120.5
As summarized in Table [3 PDT successfully accelerates training in the NLP domain. PDT achieves the baseline's 's best training loss 26. % faster and best validation accuracy 20.5% faster. This speedup magnitude is comparable to that observed in our vision experiments (10-40%). PDT also achieves higher validation accuracy of (88.44% compared to the baseline's 's 86. 10%). Fig. presents the training dynamics comparison between baseline SGD and PDT on the AG News dataset.
The successful application of PDT text classification with LSTMs provides strong evidence for its cross-domain generalizability. Despite the different optimization landscape of LSTMs, PDT achieves consistent speedups across both CNN-based vision tasks and RNN-based language tasks.
32

---

Published as a conference paper at ICLR 2026
1.0
10*
1.6
Valiatton LOGS
Tean Looa
0.6
0.15
2>
A11
0.4
0.00
30
35
10
20
Eooch
Epoch
Epoch
Epoch
Figure 19: Performance comparison between baseline optimizer and PDT. Trained the AG News dataset using an LSTM model. Trained on a single Nvidia RTX A6000 GPU for 30 epochs, 1r=0.1, batch size=128, with AnnealAAne scheduler. The shaded areas represent the standard deviation across three runs with different random seeds (0, 100, 200).
33


Reviews:
[{"review_id": "4j7UDapE15", "reviewer": "Reviewer_sCUV", "summary": "The paper proposed PDT that selectively applies updates to accelerate the training. Their method can be integrated with diverse optimizers. The authors provides diverse experiments on synthetic data and realistic dataset that their methods can accelerate the training progess without sacrificing model performance.", "questions": "1. In figure 4, is the base optimizer SGD for PDT?\n2. In line 1041, the citation is missing: \"The PDT-related hyperparameters mentioned in Sec. ?? \".", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Their method can be integrated with lots of classic optimizers as a plug-in tool.\n2. The seletive design ensures the stability of training\n3. The predictive power of PDT accelerates the convergence (in iteration-wise sense).", "weaknesses": "1. The measurement of runtime is not well explained. I'm not sure if the total runtime includes the prediction time for accelerated training for PDT. \n2. The training results cannot match SOTA sometimes. For example, in figure 10d, the training on ImageNet-1K using ViT-Base can achieve ~76% validation accuracy but in the figure the best is below 70%. Similar for the training of ImageNet-1K on ResNet50. \n3. The complexity analysis didn't cover the space analysis. So I'm not sure if it can extends to large-scale trainings as in these trainings space is also very critical.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RkzVlp1wj9", "reviewer": "Reviewer_efcL", "summary": "This paper leverages the predictive capability of KOT and introduces a framework PDT to accelerate neural network training by allowing weight updates to bypass standard gradient steps and use predicted weights instead. A dynamic consistency analysis is proposed to ensure stability when scaling to larger networks, and experiments are conducted to demonstrate empirical gains.", "questions": "- Is $\\mathcal{K}=\\sum_k \\lambda_k \\phi_k $? Is $c_k = \\langle \\phi_k, g\\rangle$?\n- What is data matrices $W_i$ and $W_{i+1}$ in neural networks? Are they neural network weights parameters?\n- When you say $A$ can be solved but you actually mean approximate? cuz equation (4) is not equal\n- how frequent you apply the prediciton step? does that depends on the optimziers, network achitecures, datasets?\n\n### **Questions for the Authors**\n\n- Is $\\mathcal{K}=\\sum_k \\lambda_k \\phi_k $? Is $c_k = \\langle \\phi_k, g\\rangle$?\n\n- What are the data matrices $W_i$ and $W_{i+1}$ in the context of neural networks? Are they weight parameters?\n\n- How frequently is the Koopman-based prediction step applied during training? Is the prediction frequency fixed or adaptive? Does it depend on the choice of optimizer, network architecture, or dataset?\n\n- How is the **binary mask** implemented in practice? Is it computed at every prediction?  \n\n- What is the computational cost of performing the SVD at each prediction step? How large are the snapshot matrices used in your experiments?\n\n- How sensitive is PDT to the number of snapshots used for Koopman estimation? Have you tested different window lengths or ranks for the decomposition?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "1. The paper explores an interesting intersection between optimal control and DNNs training, highlighting a potentially useful connection between the two fields.\n2. The authors conduct experiments across multiple optimizers, architectures, and datasets to demonstrate the generality and potential improvements of the proposed method.", "weaknesses": "- The paper’s writing and logical flow make it difficult to understand what the proposed method actually does. The main text introduces several abstract notations (\\(x_i, w_i, W_i, g, K, T\\)) with minimal explanation or connection to gradient-based optimization. For example, the authors use \\(x\\), \\(w\\), and \\(W\\) all as inputs to \\(g\\), which easily confuses readers. The overall presentation would greatly benefit from clearer definitions, consistent notation, and explicit mapping to standard training dynamics.\n\n- Most essential implementation details about how the Koopman operator is estimated, how prediction integrates with standard optimizers, and how the dynamic consistency mask is computed—are deferred to the appendix. Since reading appendices is optional during review, this omission makes it impossible to fully evaluate the method’s soundness, novelty, and overall contribution from the main paper alone.\n\n- The toy example provided does not effectively illustrate why or how the method applies to general neural network training. Its formulation lacks connection even to a simple MLP and therefore fails to provide intuition or justification for why Koopman-based prediction should improve training efficiency.\n\n- The paper mentions that predicted updates are applied \"once a while,\" but it does not specifies or studies, either theoretically or empirically, how often these predictions should be used. This frequency is critical for both computational cost and model performance. Without an ablation study on this factor, it is difficult to assess the reliability of the reported improvements.\n\n- The \"mask,\" which seems to determine when to apply predicted updates, is central to the method’s stability but is not properly introduced or explained in the main text. It is defined only in the appendix as a binary gate based on heuristics. The paper does not justify why a mask is necessary, why it must be binary, or how this design choice impacts performance or stability.\n\n- The method requires repeated SVD to estimate the Koopman operator, yet the paper provides no analysis of the associated computational overhead. There is no information about matrix sizes, decomposition frequency, or runtime breakdown. Since SVD can have cubic complexity, this omission raises serious concerns about scalability—especially given the claim of “10–40% faster convergence.”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mnxM6zfogr", "reviewer": "Reviewer_SwkT", "summary": "This paper addresses the problem of accelerating neural network training through predictive differential training - that is, predicting the future weight values of the network during training to bypass gradient descent steps. To predict the weight values at future epochs, the authors apply dynamic mode decomposition (DMD) over a small window of past weight values. The authors then show that DMD alone is insufficient for stable training, and propose masking strategy that determines which parameters will be accelerated via DMD, and which will use SGD-based updates. Through extensive experiments, the authors demonstrate that their algorithm can enhance many different types of optimizers. Also the proposed algorithm is found to work on both image classification and self-supervised learning, further highlighting the robustness of the study.", "questions": "1.\tWhile the authors have adequately provided information on the computational complexity of their algorithm, I am also curious about the memory complexity. Naively thinking, I imagine that the additional memory required would scale as $N*h$ and then some more, related to the cost of performing SVD on a $N\\times h$ matrix. Can the authors provide a comparison between their algorithm and SGD on this front? Does this additional memory requirement present any difficulties, or is this not an issue in practice?\n\n2. For the proposed masking strategy, I am curious about the relative importance between the acceleration effectiveness criterion and the dynamic consistency criterion. Can the authors present an ablation experiment by using only one or the other? I am also curious if the changes in the masking ratio during training (A.7) is related to the different contributions of the two masking criteria during different stages of training.\n\n3. The authors do mention that their method comes with couple of additional parameters and discusses their effect in A.9. Regarding this point, can the authors provide a rule of thumb for selecting these values given a training problem?\n\n4. I am confused about the direction of the inequality in the acceleration criterion. While the main text seems so suggest that equation (6) is correct, the direction of the inequality between equation (6) and (8) is different. Can the authors clarify?\n\n5. For each of the experiments, can the authors provide information of how many batches were in an epoch? I wan to get a better idea of how many gradient steps are being skipped by PDT method.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1.\tThe proposed algorithm has multiple well motivated components, that ultimately lead to significant speed ups across multiple architectures and learning tasks. The study is not a simple application of DMD either, as the authors demonstrate that their masking strategy is crucial in speeding up model training. \n2.\tThe proposed method seems to be quite robust, showing performance improvement across not just different architectures, but also completely different problem types (self-supervised learning), optimizers, and out-of-distribution minibatch situations (A.8). These results are quite impressive.\n3.\tThe paper is very clear, with the motivation and details of the algorithm clearly laid out. The use of toy problem also helps intuition and the authors provide extensive information about their experiments in the appendix, making their paper quite transparent.", "weaknesses": "1.\tThe convergence analysis in A.5 seems only to be a sketch of a proof and not a proper convergence proof. Furthermore, the inequality direction in equation 8 is different from the one show in the main text (equation 6).\n2.\tThe proposed PDT method introduces 4 additional hyperparameters, which is a bit too many to select via brute force hyperparameter search. A heuristic for determining these values will be handy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0Rt4UbESJO", "reviewer": "Reviewer_JHov", "summary": "*disclaimer* I only used LLM for understanding of part of sec 3.2. \n\nThe paper proposes an acceleration strategy for accelerating first order methods (SGD and variants) esp. in large scale deep learning systems. The strategy is to use linear approximation to approximate the nonlinear dynamics of weight updates in DNN with first order methods. In particular, the authors propose a predictive differential training  which serves as plug in tool in SGD or variant optimizer. It starts with SGD updates for a few steps and then use these snapshots to linear approximate and forecast a few steps in the future.  To accept the forecasted weights one by one, the author proposed the two criteria: acceleration effectiveness criterion and dynamic consistency criterion so that the predicted dynamics can be in accordance to the global optimizer dynamics. Empirical results show the effective acceleration on large scale systems in Vision applications and networks.", "questions": "1 In Line 371, the paper states: “the masked ratio always starts with higher values in the early stage of the training process, then generally decreases as training progresses.” The authors further suggest that this reflects the increasing complexity of training dynamics in large networks on large datasets, making them harder to predict. I am not entirely convinced by this explanation. \nIf I understand correctly, the weight is masked  means the predicted weight is accept. \nIntuitively, this trend may also be influenced by how stochastic gradient descent (and its variants) behaves: early in training, the loss landscape is typically easier to optimize, leading to faster reduction in loss and more stable gradient directions, which in turn may allow more weights to pass the masking criteria. Later in training, as the optimizer approaches (local) minima, gradients oscillate more around the optimum, making it harder to accept predicted weights for prediction horizons like 5 steps. Could the authors clarify or comment on whether this trend may be more related to optimizer behavior rather than intrinsic “complexity of the dynamics”?\n\n2. Has the authors conducted experiments analyzing the impact of different initial learning rates on the masked ratio behavior? Intuitively, a smaller learning rate might lead to more stable gradient directions, thereby increasing the proportion of weights that satisfy the masking criteria early in training. Conversely, larger initial learning rates might introduce more variability and lower the acceptance ratio. That also affects the computational analysis and efficiency of the current approach in the appendix.  \n\n3. The current empirical evaluation focuses exclusively on vision tasks (e.g., AlexNet, ResNet, ViT). While these results are compelling, it would significantly strengthen the work if the method were also evaluated on other domains, such as NLP or speech models. This would help demonstrate the generality and robustness of the PDT framework across architectures and modalities.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strengths include 1) theoretical inspired by koopman operator theory . The authors adapts to a computational efficient approach using linear approximation. (and some numerical decomposition for predicted weights) 2) the authors also show the justification of using masking strategy to make “Acceleration” possible and meaningful. 3). The quality of experiments are very high by not only showing the loss and acc reduction but also computation analysis, and show a variety of 1st order optimizers . 3) the topic is interesting in DL optimization community. 4) the presention of this paper is always good to follow.", "weaknesses": "Weakness include 1).  No thorough literature search in this line of work: for example a) Introspection:accelerating neural network training by learning weight evolution. b) Learning to Boost Training by Periodic Nowcasting Near Future Weights and a few other following this line of work.  Correspondingly, authors should include some baselines from this line of work for highlighting the superiority of proposed method  2) The proof of convergence should be discussed in more detail for theorem 1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)