Neural ODE Transformers: Analyzing Internal Dynamics and Adaptive Fine-tuning
Authors: Anh Tong
Thanh Nguyen-Tang
Dongeun Lee
Duc Nguyen
Toan Tran
David Leo Wright Hall
Cheongwoong Kang
Jaesik Choi
Venue: ICLR 2025 Poster (2025)
Keywords: Neural ODEs
Transformers
Adaptive finetune
Primary area: foundation or frontier models, including LLMs
TL;DR: This paper introduces a new transformer architecture using non-autonomous neural ODEs that allows for adaptive fine-tuning while analyzing internal model dynamics through spectral analysis.

Abstract: Recent advancements in large language models (LLMs) based on transformer architectures have sparked significant interest in understanding their inner workings. In this paper, we introduce a novel approach to modeling transformer architectures using highly flexible non-autonomous neural ordinary differential equations (ODEs). Our proposed model parameterizes all weights of attention and feed-forward blocks through neural networks, expressing these weights as functions of a continuous layer index. Through spectral analysis of the model's dynamics, we uncover an increase in eigenvalue magnitude that challenges the weight-sharing assumption prevalent in existing theoretical studies. We also leverage the Lyapunov exponent to examine token-level sensitivity, enhancing model interpretability. Our neural ODE transformer demonstrates performance comparable to or better than vanilla transformers across various configurations and datasets, while offering flexible fine-tuning capabilities that can adapt to different architectural constraints.

Published as a conference paper at ICLR 2025
NEURAL ODE TRANSFORMERS: ANALYZING INTER- NAL DYNAMICS AND ADAPTIVE FININGUNING
Anh Tong
Thanh Nguyen-Tang
Korea University anttoong12@koreaaaacc ac.
Johns Hopkins University nguyent@cs jhu edu
Dongeun Lee
Duc Nguyen
Texas A&M University-Commene dongeun, Lee@tammen edu
Qualcomm Al Research
ducanguy@qtic qualcomm. com
Toan Tran
David Hall Stanford
Cheongwoong Kang
Qualcomm Al Research*
KAIST Cw www.kangkaiist.cc.o ac..
ttrangqti tww.luualooooooooo conm. com
dLwh@stanford.edd
Jaesik Choi
KAIST, INEEJT
jaesik. wwwi@kaist.aac.ooo ac.kr
ABSTRACT
Recent udvancements in large language models (LLMs) based on transformer chitectures have sparked significant interest in understanding their inner workings. In this paper, introduce novel approach to modeling transformer architec- tures using highly flexible non-autonomous neural ordinary differential equations (ODEs). Our proposed model parameterizes all weights of attention and feed- forward blocks through neural networks, expressing these weights as functions of continuous layer index. Through spectral analysis of the model's dynam- ics, we uncover an increase in eigenvalue magnitude that challenges the weight- sharing assumption prevalent in existing theoretical studies. We also leverage the Lyapunov exponent to examine token-level sensitivity, enhancing model inter- pretability. Our neural ODE transformer demonstrates performance comparable to or better than vanilla transformers across various configurations and datasets, while offering flexible fine-tuning capabilities that can adapt to different architec- tural constraints.
INTRODUCTION
The recent advancements in large language models (LLMs) based transformer architec- tures (Vaswani et al., 2017; Brown et al., 2020) have impressive empirical performance in various tasks. These transformer models with millions billions of parameters, pose a formidable chal- lenge in comprehending the intricate processes transpiring within their layers. Initial efforts have been undertaken unravel mechanistic interpretability (Elhage al., 2021) and adopt mathemati- cal approach (Geshkovski et al., 2023a;b) to better understand the inner workings of these complex models.
As the field LLMs continues to evolve and grow, is important to develop fundamental approach to understanding such complex models. This can provide stepping stones improving trustwor- thiness, interpretability, and alignment in LLMSs (Shevlane et al., 2023). To this aim, there are multifaceted approaches combining insights from mechanistic interpretability. mathematical analy- sis, and other complementary methods. Rigorous approaches from Geshkovski et al. (2023a;b) open
QQalcomm Vietnam Company Limited

---

Published as a conference paper at ICLR 2025
interesting directions, providing theoretical understanding of transformers using neural ordinary dif- ferential equations (neural ODEs) (Chen et al., 2018). In practical applications, neural ODEs have been used in models proposed by Lu et al. (2019); Zhong et al. (2022); Dutta al. (2021). While Geshkovski et al. (2023a;b) suggest that the model outputs collapse to small number of points at infinite depth, Zhong al. (2022); Dutta et al. (2021) demonstrate the cases of neural ODE- based transformers. However. the main limitation in this line of work is that they often rely on shared- weight assumptions, which leads to a discrepancy between theoretical results and practical observations. Theoretically, the clustering behavior of fixed points in these studies (Geshkovski al., 2023a;b) does not align well with the outputs of autoregressive transformers, which flexibly predict next tokens rather than forming clusters. Practically, the potential of differential equation- like adaptive computation often overlooked, as evidenced by the limited exploration of neural ODE versions of transformers with time-dependent vector fields, as opposed to weight-sharing ap- proaches.
In this paper, introduce novel approach to modeling transformer architectures using highly flex- ible non-autonomous neural ODEs'. Apart from existing methodologies that seek weight-sharing, proposed model generates all the weights of attention and feed-forward blocks through hyper- neural networks. Specifically, we consider weights as functions of continuous layer index (or time). We conduct comprehensive exploration and analysis of its dynamies through the lens of dy- namical systems. Our primary analytical tool is spectral analysis that examines the components of our models including attention blocks and feed-forward blocks. Our analysis reveals some important properties of our models. The observed spectral dynamics tend to exhibit increasing magnitudes of cigenvalues. Through simulation, we demonstrate that this phenomenon inhibits the emergence of clusters, contrasting with the results reported in Geshkovski et al. (2023a). In other aspects, we use the Lyapunov exponent to study token-level sensitivity, which can be a potential tool for explainable machine learning methods (Gunning al., 2019).
Furthermore, we present novel approach that leverages the adaptive computation of differential equations to create a versatile pretrained model. Our model exhibits a unique characteristic: the abil- ity to be fine-tuned under various architectures derived from the initial pretrained structure. When presented with novel datasets, our model dynamically adjusts the step size of its ODE solvers. This adaptation effectively transform the model's architecture, allowing it to function as a vanilla trans- former. Consequently, this enables the application of wide range of fine-tuning techniques such as LoRA (Hu et al., 2021). To the best of our knowledge, this represents the first instance a pretrained model capable of being fine-tuned across diverse architectural configurations.
The contributions of this paper are threefold. First, we introduce neural ODE transformers having comparable performance with GPT in different settings and data sets. Second, offer a compre- hensive analysis of the proposed models, unraveling their intricate internal dynamics and a tech- nique to assess foken-level sensitivity, enhancing our ability interpret model behaviors. Lastly, we show that our model provides flexible computation fine-tune tune while maintaining competitive performance.
REVISITING TRANSFORMER
Transformers proposed by Vaswani et al. (2017) consist of key components, including self-attention and feed-forward neural network.
Attention mechanism Consider sequence with n tokens, X1, X...... X,, in Rdd, and represent it as matrix X E Rexd, Let us define three linear projections Q, K, E Rddmxd for query, key, and value, respectively. The self-attention operation expressed as follows:
[QQ [XQTXX' KXT
Autn(Q, K, V) = softmax
(XVT).
(I)
Vd
In essence, the self-attention mechanism acts as operator determining where to focus attention within the entire sequence. The softmax function introduces emphasis, making it more likely that one token will be highlighted over others.
Non-autonomous ODEs have time-dependent vector fields.
2

---

Published as a conference paper at ICLR 2025
Self-attention is further extended into multi- melti-headdee attention. In this configuration, multiple indepen- dent self-attention outputs, referred to as heads, are concatenated and reprojected, allowing for the introduction of residual connections. The formal definition is given by:
MultiHead (Concat(head)) headyy )O, where head, :=ttt (Q,, K,, X); = 1, ,,....
Here, H is the number of heads and O E Rdowe x is projection matrix.
Feed-forward block The second crucial component is the feed-forward layer, modeled as a two- layer multilayer perceptron with nonlinear activation function, expressed as:
GeLU(XW, + W b2.
where Wi E RRdipxd W2 Rdxdur, bj E Rdmv, and b2 E RR.. The feed-forward layer, singular component in the transformer, introduces nonlinearity in the computational flow, enhancing the model ability capture complex patterns and relationships.
RELATED WORK
The work in Lu et al. (2019) is pioneering in establishing the link between neural ODEs and trans- formers. They recognize that the residual connections in transformer blocks can be interpreted as a Lie-Trotter splitting scheme akin ODE solvers. Building this insight, Lu et al. (2019) proposes a novel architecture called MacaronNet, motivated by an enhanced numerical solver. In a parallel vein, Li al. (2021) presents a new variant of transformers inspired by higher-order solvers to address machine translation tasks. Our approach differs from these works in key aspect: while they were inspired by ODE solvers to develop novel transformer construction methods, focus on directly formulating differential equations. Consequently, our model enjoys all the properties of neural ODEs, such as adaptive computation, and is compatible with any ODE solver.
The approach outlined in Dutta et al. (2021) aligns with certain aspects of our proposed model. Specifically, this approach involves introducing time-evolving attention layers and feed-forward lay- ers. In the context of attention mechanisms, Dutta al. (2021) suggests concatenating inputs and layer embeddings. However, it worth noting that this approach may need more flexibility, as there persists shared component in the computation of attention weights. This shared component relies solely on input and remains independent of layer embeddings.
In another investigation, Zhong et al. (2022) explores various configurations of weight-sharing strategies when training transformers in a neural ODE style. The key finding indicates that as the degree of weight-sharing between layers increases, the model's performance tends to degrade. Our model overcomes this limitation by using time-dependent weights, resulting in improved perfor- mance.
body of research has emerged, focusing on constructing models and laying down theoretical frameworks for transformers built upon neural ODEs. works such as Geshkovski al. (2023a;b), transformers are conceptualized as particle systems, and the geometric behaviors of simplified equiv- alents based on neural ODEs analyzed. The predominant findings in these studies highlight the tendency of systems to converge to clusters as they approach their limits. Interestingly, Sink- former (Sander et al., 2022) also uses similar theoretical approach but aims design new vari- ant of transformers. However, we note that the weight-sharing assumption of Geshkovski et al. (2023a;b) is rather restrictive. Our work demonstrates that learned weight dynamics prevent mod- els from suffering clustering effects. This suggests reconsider alternative assumptions for such studies.
Research has been devoted to deep equilibrium models, as demonstrated by Bai et al. (2019), show- casing a close relationship with neural ODEs. This technique has been effectively employed in training transformer-lis like models with unspecified depth during training. While similar in some spects, our model (and neural ODEs in general) allows for adaptive fine-tuning with flexible ODE solver step sizes.
3

---

Published as a conference paper at ICLR 2025
(b)
Linear
GeLU
K
Temporal Embedding
MLP
Proj
MatMul
W (t)
Linear
duehh
dass
W(()
Tamporal Embedding
Figure (a) The vector field of DIFFEQFORMER with attention block and feed-forward block constructed from dime-dependent weights. (b) The architecture of time-dependent weights.
TIME-DEPENDENT WEIGHT TRANSFORMER
This section introduces our proposed transformers as differential equations, called DIRFEOFORMER. In the beginning, we formulate the transformer within the context of ODEs. We then provide the approach for representing time-dependent weights in our models.
4.1 TRANSFORMERS THROOGH THE LENS OF DIFFERENTIAL EQUATIONS
Consider sequence XI...., X,, with X; E R4. Let dynamical system evolve the given sequence under the following differential equation:
(t) f(; 2 (t, ()),
E [0, 7], 1, .........
5;, (0) X,;,
(3)
Here, (tx(t)= = dz; and Zmn (t) denotes [z] (t). In (t)]. This formulation was first presented in Lu di
al. (2019). Based numerical solver, they proposed the MacaronNet architecture with discrete layers, in contrast to our approach which uses continuous layers with time time-dependent weights (see Section 4.2 time-dependent weights).
Here f acts as an attention mechanism, capturing inter-particle interactions. Self-attention facilitates the transfer of information among particles. In a multi-heat setting, each particle receives informa- tion from H. which is the number of different versions of all particles, and updates itself based on this information. express this, we define f as follows:
H Σ
(())))))) vi(t)a,(t), Vd
(())
(4)
Aml LiA
where
(x(.), (-),), V.(+) are the mapping from [0, T] to Rdmxd
Lin
E exp (z, is the normalizing term in the softmax function and
(r, = y.
On the other hand, the convection g serves a role as feed-forward blocks having all time-dependent weights:
g(m,(1),t) W (I)GeLU(WW* (1) (t).
Figure la illustrates how construct the vector field of DIFFEOFORMER.
Remarks The particle system exhibits certain similarities with consensus-bel bassed optimization tech- niques, explored in previous works (Pinnau al., 2017; Carrillo et al., 2021), with initial
4

---

Published as a conference paper at ICLR 2025
discussion provided in Geshkovski al. (2023b). Drawing upon this resemblance, it becomes pos- sible to elucidate specific observations in transformers, revealing that these models implicitly en- gage in optimization processes mnesa-optimization)) over contextual information (Garg et al., 2022; Von Oswald et al., 2023; von Oswald et al., 2023).
Notably, the work of Elhage et al. (2021) introduces the concept of residual streams: the connec- tions between input tokens and their corresponding output tokens. In this framework, the role of attention facilitate the communication and information transfer between these residual streams. Multi-head attention encompasses information from various subspaces, allowing for the merging and writing of information residual streams through residual connections. This perspective aligns with the interpretation of particle systems described in Lu et al. (2019); Geshkovski et al. (2023a;b).
Additionally, placing attention layers and feed-forward layers at the same level was investigated in prior studies (Peng et al., 2022; Zhong et al., 2022; Dehghani et al., 2023). This approach allows us fully conceptualize transformers as neural ODEs with Euler methods, in contrast the Lie- ie-Trotter scheme employed in Lu et al. (2019).
4.2 REPRESENTATION
IME-DEPENDEET WEIGITS
We model time-dependent weights for attention components (4) (t),,, (t), and feed-forward weights W (1), (t), using time-dependent unit W 14, + di Xdwa that embeds the time information in the Fourier domain, inspired by the existing diffusion model architectures (Song et al., 2020; Peebles & Xie, 2023). In particular, for any time E R+,
W(t) = Proj (Sinusoiati)i
(6)
where Sinusoidal R -> didinii is used to embed time t into a higher-dimensional space and the two-layer multilayer perceptron MLP Radusssiaal -> Re is used. The MLP's output will then be linearly projected and reshaped to match the desire of the weight matrices by using the operator Proj: dR RRd xd... Figure 1b depicts the architecture of W
Remark The time-dependent weights can be understood as hypernetworks (Ha et al., 2016; Stanley al., 2009). However, our approach stems from distinct motivation, aiming establish flexible method for modeling time-dependent vector fields in differential equations, employing a Fourier time embedding representation rather than the architecture proposed in Ha et al. (2016). The frame- work presented by Ha et al. (2016) contains two types of hyperetworks: static and dynamic. Our approach aligns more closely with the static variant. Conversely, existing work, such as Zhang et al. (2019); Choromanski al. (2020), in the realm of neural ODE, attempts to parameterize weights using differential equations, resembling a form of dynamic hypernetworks. While this paper does not explore this approach, is noteworthy that the main challenge lies in considering the state of the original ODE, which is high- high-dimensional and could result in computational expense.
5 ANALYZING INTERNAL DYNAMICS OF TRANSFORMERS
In this section, we explore the qualitative aspects and properties of our proposed model, including its behaviors of attention mechanism and model interpretability. Our analysis presented here is conducted on the EOFRRMER model trained on the OPEN dataset (refer to Section 6 for the detailed setup).
5.1 ANALYSIS OF ATTENTION MECHANISM
To characterize attention blocks in DIRFEQFORMER, we will use existing tools, including gradient flows (Geshkovski et al., 2023a;b) and induction heads (Elhage et al., 2021).
Analyzing Query-Key (QK) pair First, explore the spectral flow of the Query-Key (QK) pair, Qu(t) (T,t), This particular element plays a vital role in determining particle communication and provides insights into the extent of their interactions.
Figure 2a displays the dynamics of eigenvalues of (t)))) (t) in DIFFEOFORMER model. exhibits higher degree of interactions in its last layers, evidenced by the broader range of
5

---

Published as a conference paper at ICLR 2025
Read
Head
Head 2
Head 3
Hend
Head
Head
Head
Realin
Layar
Layer
Layer
Layer
Lager
Layer
Layer
Layer
(a) Spectral dynamic of QK
(b) Spectral dynamic of OV
(c)
(d)
(e)
Figure 2: (a-b): Spectral dynamics of QK and OV pairs. (c-d): Trajectory of sequence consist- ing of 40 points 3-dimensional space. (c) Attention-only model with shared weight assump- tions as described in Geshkovski et al. (2023a): Clusters emerge. (d) Attention-only model with time-dependent weights of increasing magnitude, inspired by observations from trained DIFFEQ- FORMER: No clusters occur. (e) Plot of function in simulation that mimics the magnitude of Q(t), (t),, (t) over time like in trained DORMER.
eigenvalues in these layers. The eigenvalues of (!) (L) directly affect the variance of I; (t)Qx(t) KA (t). For instance, if 2, ~ NO, I). then Var(æ: Ay) trace(A A A) >, (AA) where (A(A) is the i-th eigenvalue of A. A high variance of x (t) Kn (t)xz I,....
indicates that the softmax of [x. (1) K₁ (()=)((()))))))))))))))))))))( (Q..=,,...... tends to approach one-hot vector. i-th particles.
This results sharper attention weights, leading to more distinct exchange between the j-th and
To gain deeper understanding, we make connection with the work of Geshkovski al. (2023b). In this study, the authors simplify the model by assuming that Qtt), ((t), and (tt) are all identity matrices, i.e.
1 Σ (t), Lg
(,(2,(t),
where Lg Σ exp(B(x;;()) (0), Here, the parameter 3 serves as simplification of Qu(t)TKa,(t) (t) in equation 4. The study by Geshkovski et al. (2023b) looks into the interaction energy, and its gradient flows govern the dynamics of fa (Otto, 2001; Jordan et al., 1998; Ambro- sio et al., 2005; Villani al., 2009). Importantly, the parameter 3 directly influences the speed at which the interaction energy decreases or increases. Such interactions in DIFFEQFORMER can be reflected by the eigenvalues (t) (t). For example, a wide range of cigenvalues encourage interactions between particles, as they indicate stronger coupling and information exchange within the system.
Analyzing Output-Valee (OV) pair While ( identifies the source and destination particles of inter-partie communication. the Output- (OV) pair represents the component describing how information of the source particle is merged to the destination particle.
Figure 2b illustrates the spectral dynamics of the OV matrix in DIFFEQFORMER. A significant number of heads exhibit positive eigenvalues, which are primarily concentrated near the output layer. It seems to have correlation with the spectral dynamics of the QK pair.
Consider a communication from a source particle j to a destination particle The OV matrix has eigenvalues (Ae )d and corresponding eigenvectors (vx)==1- We express the states of

---

Published as a conference paper at ICLR 2025
particles i and j at time t as linear combinations of these cigenvectors:
(t)
Wk Uks
2y(t)
w
Uk:
where {wt Vi=z and fug k= are the respective coefficients. The output of the Euler step for atten- tion can be interpreted as:
d
xz + A1)= Σ (Aew Al +
Here, focus the coefficient A AL wt- The magnitude of d determines the influence of the source particle (1) on the output (1 At). Larger absolute values of X result in a stronger influence of the source particle on the destination particle's next state. The sign of λ decides either positive or negative influence of source particle with respect to the corresponding eigenvector ey.
Drawing connection to the work of Geshkovski et al. (2023a:b), the authors suggest two simplified cases: the attractive scenario, where particles exhibit attractions and more likely to follow some leader particles with V = Igi and the repulsive scenario with V -Id. where particles repel each other. We extend this analysis by examining both the sign and magnitude of eigenvalues A as shown in Figure 2a and 2b. This means that DIFFEOFORMER shows both attractive and repulsive behaviors, with attraction being more dominant in some attention heads near the last layers.
In comparison, the approach presented by Elhage et al. (2021) characterizes OV pairs as exhibit- ing copying behavior when the eigenvalues of the OV pair feature numerous positive values. This copying behavior conceptually is similar to the attractive scenario in Geshkovski et al. (2023a;b). In mathematical terms, V Is means that all eigenvalues are positive and equal to 1. If copying behaviors occur frequently, it indicates a selective focus on specific particles, resembling a situation where only subset of major particles is highlighted. Note that copying behavior is characterized as a component of induction heads, recognized as a key observation associated with in-context learning performance in Elhage et al. (2021); Olsson et al. (2022).
Under the interpretation of Elhage et al. (2021), induction heads perform two key operations: (i) de- tecting matching patterns through QK pair and (ii) copying values when matches are found through OV pairs. The "matching" operation manifests through attention scores, where positive eigenvalues of QK pair indicate strong attention paths between similar tokens. The "copying" operation is facil- itated by the OV pair, where positive eigenvalues suggest effective information propagation. Given our observations in Figure 2a and 2b of both positive eigenvalues in QK pair (supporting pattern matching) and positive eigenvalues in the OV pair (enabling copying), it is highly probable that induction heads occur in the last layer of DIFEEOFORMER.
Clustering behavior in DIFFEOFORMER Trained DIFFEOFORMER exhibits a distinct dynamic of QK and OV pairs characterized by increasing magnitudes of eigenvalues over time, with a no- table peak observed near the final layers. This dynamic contrasts with the assumptions put forth by Geshkovski et al. (2023a), who posited the emergence of clustering behavior at limit. Given this unique dynamic, it is crucial to investigate potential clustering behavior in our DIFFEOFORMER. To this end, we simulate ODE trajectories for attention-only model with the simplified dynamic based on the trained DIFFEOFORMER our findings indicate an absence of emergent clusters for a case of time-dependent weights. This contrasts with the cluster emergence observed in the weight-sharing case, shown in Figure 2c. This outcome aligns with the practical considerations of autoregressive language modeling tasks with next-token prediction, where clustered outputs would be counterintuitive. This shows that discrep- ancy between the theoretical assumptions in existing studies and the empirical properties exhibited by our DIFRMRR models. Therefore, it necessary to have future rigorous investigation into this gap.
(refer to Appendix C.8 for details). As illustrated in Figure 2d,
Remark Due to the continuity of weights within our model, the dynamics of eigenvalues exhibit the continuity property, as discussed in previous works (Wilkinson, 1965; Hoffman & Wielandt, 2003). Consequently, a smooth transition of spectral information between layers can be observed in our analysis of QK and OV pairs. In contrast, the spectral flow in the case of vanilla transformers do not exhibit clear patterns or behaviors in their dynamics, making it challenging to identify induction heads (see Appendix E).
7

---

Published as a conference paper at ICLR 2025
BIGBANG one of those musical entities that transcends langunge. It's one of those rure groups that both innovates and defines the direction genre takes, Covering sound that inclades hip hop, R&B and electronic dance, BIGBANG and
SEOUL (Reuters) South Korean Lee Sedal won his first match against computer program devekoped by a Google subsidiary on Sunday in the ancient hoard game Gin denying clean sweep for the artificial intelligence in five-mateh series. The world's top player Lee Sedol reviews the match after the fourth
(b)
Figure 3:
Lyapunov exponent values represent the sensitivity of previous words to the next word. Higher values correspond to more intense highlighting in red. (a) The next word is its. (b) The next word is match.
5.2 SENSITIVITY ANALYSIS WITH LYAPUNOV EXPONENTS
Next, we utilize the Lyapunov exponent assess the sensitivity of the inputs of DIFFEOFORMER. Lyapunov exponents can provide valuable insights into the behavior of transformer models by quan- tifying how small changes in the input may impact the outputs. Previous studies have utilized Lyapunov exponents for understanding neural networks, primarily focusing recurrent neural networks (RNNs) (Vogt et al., 2022; Storm et al., 2023; Engelken et al., 2023; Engelken, 2023). In another example, Gilpin (2021) demonstrates how Lyapunov exponents enhance the inter- pretability of time series data (see Appendix for a brief review).
The main challenge in extracting Lyapunov exponents for the entire DIFFEOFORMER system stems from the high dimensionality of its states. Calculating Lyapunov exponents requires computing the spectrum QR decomposition of Jacobian matrices, which can be computationally infeasible for DIFFEQFORMER
In this context, our objective is investigate the token- level sensitivity of transformers. Specifically, we aim to understand how the changes of the i-th token in a given input sentence influence the j-th token of its corresponding output. To achieve this, intend monitor the evolution of the tangent vectors associated with the i-th input particle and j-th output particle, similar to the formulation in equation 8, using the following ODE:
(1 = Jos(0))
with (0) = Ia and = field equation 4. The Jacobian Jis comprises two components: the attention information and the change between input and output. Consequently, employing Lyapunov exponents can provide a more principled approach to understanding the behavior of transformer models compared solely using attention matrices, which have limitations as discussed by Jain & Wallace (2019) (see further discussion in Appendix B).
=
+ (t) is the Jacobian of the vector
Our approach may share similarities with conditional Lyapunov exponents, which analyze the sta- bility of specific subsystems within larger system (Pecora et al., 1997). Like these exponents, our sensitivity values provide insights into the stability of specific tokens within the transformer model.
Figure 3 showcases how our approach quantifies sensitivity for predicting the next words using Lyapunov exponents. In these examples, the obtained Lyapunov exponents can explain the case of possessive pronouns. For instance, in the BIGBANG example, the next word "its" corresponds to "BIGBANG", which is highlighted by the Lyapunov exponent. Similarly, in the Alpha Go arti- cle example, the obtained Lyapunov exponent can highlight the next word "match" and the order preposition "after." This analysis offers potential avenues for interpreting the inner workings of transformers, building upon prior efforts in attention interpretability (Jain Wallace, 2019; Wiegr- effe & Pinter, 2019; Chefer al., 2021; Ali al., 2022). Further results and examples of this analysis, including comparisons with attention explanations can be found in Appendix F.
6 EXPERIMENTAL RESULTS
This section presents an empirical comparison of our proposed model DIFFEOFORMER and vari- ous baseline models in the context of language modeling tasks. The implementation of our model
For example, consider model taking 1000 tokens with an embedding dimension of 768 as input; the dimension of its would be 768, 000.
8

---

Published as a conference paper at ICLR 2025
WIKITEXTI03
OPENWEBTEXT
VMI parp MESTT
208 600 DER
300 1838
180 038 500
KX
IRR
1390 1633 20m
Stop
(a) small
(b) medium
(c) small
(d) medium
(c) large
Figure 4: Validation Perplexity of DIFFFOFoRMER in comparison with GPT models. (a,b) Results on WIKITEXTIO 03 dataset two architecture settings. (c, d, e) Results on OPENWEBTEX dataset on three architecture settings.
and baseline models built on JAX (Bradbury et al., 2018), utilizing an ecosystem that includes Equinox (Kidger & Garcia, 2021), Haliax, and the Levanter framework (Hall al., 2023). Our implementation available at https: https://gittuttt //github. wwt.coo////////ee SDML-DU/ /KV//lw..
6.1 LANGUAGE MODELING
We evaluate our model DIFFEOFORMER, in comparison baselines both OPENWEB- TEXT (Gokaslan et al., 2019) and WIKITEXT103 (Merity et al., 2016) for the autoregressive mod- eling task. The main evaluation metric is perplexity, a commonly used measure this task.
Baselines We explore various baselines, with the primary one being decoder-only transformers in- cluding GPT (Radford, 2018; Brown et al., 2020) and Llama (Touvron al., 2023). We extend our investigation to broader range of baselines for the Wikitext dataset, incorporating techniques such as weight-sharing (Zhong et al., 2022). However, on the OPENWEB TEXT dataset, we exclusively compare model against the GPT model.
Setting We consider model settings including: GPT-medium, GPT-small, GPT-large, and Llama- 1B (see Table in Appendix for details). We ensure uniformity in terms of model configuration between DO and baselines, including the depth of the models, the dimension of token embeddings, the number of heads, and other hyperparameters.
Table 2: Perplexity on OPENWEB dataset.
Table 1: Perplexity on WIKITEXT 103 data set.
Model
Perplexity
Model
Perplexity
GPT-small
22.60
GPT-small
22.87
Our model-small
22.06
Shared parameter
36.27
GPT-medium
17.85
Our model-small
22.84
Our model-medium
17.67
GPT-medium
23.38*
GPT-large
15.64
Shared parameter
32.95
Our model-large
15.43
Our model-medium
21.94
Llama- 1B
10.17
Our Llama-1B
9.68
Table shows the perplexity across different models on WIKITEXTION 103 dataset. Our model achieves strong performance, outperforming GPT models while exhibiting less overfitting. Notably, GPT- medium shows signs of overfitting model parameters without a corresponding increase data (Kaplan et al., 2020). Though the shared-weight approach excels in masked language modeling (Lan et al., 2020), it falls short in autoregressive language modeling tasks this experiment.
this data (see Figure 4). This can be attributed to an increase
Table 2 demonstrates that model consistently outperforms all baselines on the OPENWEB TEXT dataset, with particularly substantial improvements over the Llama- IB baseline (see Appendix C.6).
Downstream evaluation We evaluate DITTEFFORMR against GPT-large and Llama- IB using the In-evaluationn won-aaarninnenneeeeeeeeeeeaeeeee framework (Gao et al., 2024), following the experimental protocol established in Biderman et al. (2023). Our evaluation encompasses various downstream tasks in both zero-shot and few-shot settings. While DieFEQFORMER demonstrates comparable performance to the baselines across most tasks, it achieves statistically significant inspriiiantimmoeeeeeeeee improvements reading compre-

---

Published as a conference paper at ICLR 2025
LoRA
Full runk
LoRA
Full rank
10
40
63
Oara 08 24
Vill perplesity
13
Vaa parpless poccty
Vil parpa perslexty
Vil parpp loxttt
Oars 08
Oars 08
18)
Oars (88
91
30
0P7 (24
from somit
23-
20
0:50
18
1.38 Token
26
2.58
0.35
155
1.58
2
0.56
15
1:38
25
2.58
0.38
15
1.08
26
2.08
Token
Tokenn
Token
(a) TEXT + WIKITEXTI03
(b) WIKITEXTI03 - WEBTTTT
Figure 5: Finetune validation perplexity across two settings: (a) OPENWE - WIKI- TEXT103; (b) WIKITEXT103 + WEENTE All models are pretrained with 18 function evaluations (or layers) and finetune in different settings of function evaluations 9, 12, 18, 24 with LoRA and full-rank finetune. We compare these with the baseline of corresponding GPT model pretrained and finetune with 24 layers which is much expensive computation both pretrain and finetune.
hension tasks (Lambada OpenAl and Lambada Standard) in both zero-shot and five-shot evaluations. Detailed results are presented Appendix C.7.
6.2 COMPUTATIONAL ADAPTABILITY FOR FINE-TUNE
Modeling transformers using differential equations enables adaptive computation during fine-tuning through the adjustment of ODE solver step sizes. By utilizing fixed step size, populate the weights of DORMER over time steps to obtain discrete-layer pretrained models. Now, we can apply any fine-tuning techiques such as LoRA (Hu et al., 2021).
Finetune settings The pretrained model undergoes training on one dataset before evaluation on novel data. We examine two scenarios: pretraining WIKITEXTI03 followed by fine-tuning on OPEN WEBTEXT (WIKITEXT103 + OPENWEBTEXT), and pretraining on OPEN TEXT WIKITEXT103) Both LoRA fine-
followed by fine-tuning on WIKITEXT (OPenWeRTEET) tuning and full-rank weight fine-tuning approaches are employed.
Model setting We trained DIFFEOFORMER models with 18 time steps (layers). Fine-tuning was conducted on models with varying time steps: 9, 12, 18, and 24. For instance, pretraining with 18 time steps followed by fine-tuning with 9 time steps is denoted 18 9. A baseline GPT model with 24 layers was used for both training and fine-tuning.
Result Figure 5 shows the perplexity evaluated on validation datasets. Our models demonstrate su- perior performance compared to the GPT baseline across configurations 18 + {12, 18, 24} under LoRA fine-tuning.
noteworthy that the baseline model was trained and fine-tuned with greater computational resources (24 layers). In the case of full-rank weight fine-tuning (or continued train- ing on new data), our models (18 + 24) exhibit comparable performance to the baseline, despite initial training with only 18 time steps. Furthermore, the flexibility in fine-tuning allows our models to mitigate overfitting by employing fewer time steps. This particularly evident the case of the WIKITEXT 103 dataset, where models (18 + {9, 12}) are less prone overfitting during the full-rank weight fine-tuning process.
CONCLUSION
This paper proposes novel method for formulating transformers as neural ODEs, revealing inter- esting characteristics within the models. We posit that the findings presented herein can provide theorists with foundation for developing ODE-based transformers in their theoretical frameworks. At the same time, practitioners may explore the potential of integrating ODE approaches into trans- formers. However, there remain important areas for future work such improving memory effi- ciency at larger scales, exploring advanced ODE solvers (McCallum & Foster, 2025), examining the robustness of the proposed model against adversarial attacks (Huang al., 2021) and extending to various aspects including stochastic differential equation paradigms (Tzen & Raginsky, 2019; Tong al., 2022) and alternative architectures (Tong et al., 2023).
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
The source code of this work is available at https::///// hettps://githbb cttp.ccco/ssoo/e Doom//SDMIKU//gkELL.. We built upon the Levanter library (Hall al., 2023), which manages Pseudorandom Number Generator (PRNG) states in JAX ensure bitwise determinism in both data processing and the training pro- cess. All hyperparameters are detailed in the Appendix, and corresponding configuration files can be found in the source code repository.
ACKNOWLEDGMENTS
This work was supported by Institute of Information communications Technology Planning & Evaluation (IITP) grants funded by the Korea (MVI) (No. RS-2019-11190079, Arti- ficial Intelligence Graduate School Program(Korea University); No. RSS2009--1110075, Artificial Intelligence Graduate School Program (KAIST); No. RS-2022-11220884 Development of Artificial Intelligence Technology for Personalized Plug-and-Play Explanation and Verification of Explana- tion; No. RS-2024-00457882. AI Research Hub Project) and the New Faculty Settlement Research Fund by Korea University. This work was supported by the New Faculty Settlement Research Fund by Korea University and Artificial intelligence industrial convergence cluster development project funded by the Ministry of Science and (CS(M Korea) & Gwangju Metropolitan City. We also acknowledge the Google Cloud Research Credit Program for their support during the early develop- ment of this work.
REFERENCES
Ameen Ali, Thomas Schnake, Oliver Eberle, Grégoire Montavon, Klaus-Robert Müller, and Lior Wolf. XAI for transformers: Better explanations through conservative propagation. In Interna- tional Conference on Machine Learning, pp. 435-455 451. PMLR, 2022.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the spuce of probability measures. Springer Science & Business Media, 2005.
Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. Deep equilibrium models, 2019.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brie, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
James Bradbury. Roy Frostig, Peter Hawkins. Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miin al. JAX: composable transformations of Python+ NumPy programs. 2018.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
2020.
José A Carrillo, Shi Jin, Lei Li, and Yuhua Zhu. A consensus- based global optimization method for high dimensional machine learning problems. ESAIM: Control, Optimisation and Calculus of Variations, 27:55, 2021.
Hila Chefer, Shir Gur, and Lior Wolf. Transformer interpretability beyond attention visualization. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp- 782-791, 2021.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv: 160. 2016.

---

Published as a conference paper at ICLR 2025
Krzysztof Choromanski, Jared Quiney Davis, Valerii Likhosherstov, Xingyou Song, Jean-Jacques Slotine, Jacob Varley, Honglak Lee, Adrian Weller, and Vikas Sindhwani. An ode to an ode, 2020.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 11116344-63399 2022.
George Datseris. Dynamicalsystems.jt A julia software library for chaos and nonlinear dynamics. Journal of Open Source Software, 3(23):598, mar 2018. doi: 10.21 105/joss.0059 00598. URL https: ///oi. 10. 21 2111//11..111/1111 00598.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers 22 billion parameters. In International Conference Machine Learning, pp. 7480-7512. PMLR, 2023.
Subhabrata Dutta, Tanya Gautam, Soumen Chakrabarti, and Tanmoy Chakraborty. Redesigning the transformer architecture with insights from multi-particle dynamical systems, 2021.
Nelson Elhage. Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1, 2021.
Rainer Engelken. Gradient flossing: Improving gradient descent through dynamic control of jaco- bians. arXiv preprint arXiv:2312.17300 2023.
Rainer Engelken, Fred Wolf, and L. F. Abbott. Lyapunov spectra of chaotic recurrent neural net- works. Phys. Rev. Res., 5:043044, Oct 2023. doi: 1011009PiysserRReeacch...444. URL fittps://iist aps pptpppp.rg/aaai......//////////e 1t/..011111//yyyeeeeeeee.. 5. 043044.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa. Jason Phang. Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lin- tang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024. URL https::///// org records/ 12608602.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
URL https::t///...........tt URt.httts:://ppeneee //oppeeeeeeieeeeeeeeeeeeeeeeeeee net. /aic/formne forum? i f1N.J220et..
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clus-
ters in self-attention dynamics, 2023a.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical per- spective transformers, 2023b.
William Gilpin. Chaos as an interpretable benchmark for forecasting and data-driven modelling.
In Thirty-ffthh Conference Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021. URL http::////peervvieeee net forum?id=ed jtb jtbj
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus. http: //Skylion00 g thub. wwt.p/ppeeeeeeeteerrrroooo 2019.
Audrunas Gruslys, Rémi Munos, Ivo Danihelka, Mare Lanetot, and Alex Graves. Memory-efficient backpropagation through time. Advances in neural information processing systems, 29, 2016.
David Gunning, Mark Stefik, Jaesik Choi, Timothy Miller, Simone Stumpf, and Guang-Zhong Yang. Xai exi-axplaee artificial intelligence. Science robotics, 4(37)eeay77120, 2019.
David Ha, Andrew Dai, and Quoe V. Le. Hypernetorrss, 2016.
12

---

Published as a conference paper at ICLR 2025
David Hall, Ivan Zhao, Virginia Adams, Jason Wang, and Percy Liang. Levanter, 2023. URL https://github: wttps/tdoooooodnaoeeeeoe
Alan Hoffman and Helmut W Wielandt. The variation of the spectrum of normal matrix. In Selected Papers Of Alan Hoffman: With Commentary, pp. 118-120. World Scientific, 2003.
Joachim Holzfuss and Ulrich Parlitz. Lyapunov exponents from time series. Lyapunov exponents: Proceedings, Oberwolfach, 1990, pp. 263- 1991.
Edward Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, Weizhu Chen, al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Reprexentations, 2021.
Yifei Huang, Yaodong Yu, Hongyang Zhang, Yi Ma, and Yuan Yao. Adversarial robustness of stabilized neuralodes might be from obfuscated gradients, 2021. URL abs/2003 13145.
Sarthak Jain and Byron Wallace. Attention is not explanation. arXiv preprint urXiv:190..10166 2019.
Richard Jordan, David Kinderlehrer, and Felix Outo. The variational formulation of the fokker- planck equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836/ 2020.
Patrick Kidger. On neural differential equations, 2022.
Patrick Kidger and Cristian Garcia. Equinox: neural networks in JAX via callable PyTrees and filtered transformations. Differentiable Programming workshop at Neural Information Processing Systems 2021, 2021.
Zhenzhong Lan, Mingda Chen, Sehastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. Albert: A lite bert for self-supervised learning of language representations, In International Conference on Learning Representations, 2020. URL https: http///wwwewweeeee.tteee / id=HleA7AEtvs,
Bei Li, Quan Du, Tao Zhou, Shuhan Zhou, Xin Zeng, Tong Xiao, and Jingho Zhu. Ode transformer:
An ordinary differential equation-inspired model for neural machine translation, 2021.
Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. In International Conference on Learning Representations, 2018.
Yiping Lu, Zhuohan Li, Di He, Zhiqing Sun, Bin Dong, Tao Qin, Liwei Wang, and Tie- Yan Liu. Understanding and improving transformer from a multi-particle dynamic system point of view, 2019.
Sam McCallum and James Foster. Efficient, accurate and stable gradients for neural odes, 2025. URL htttps://axxi...et abs/2410 11648.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph. Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. /185, 2022.
Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. 2001.
Louis M Pecora, Thomas L Carroll, Gregg Johnson, Douglas J Mar, and James F Heagy. Funda- mentals of synchronization chaotic systems, concepts, and applications. Chaos: An Interdisci- plinary Journal of Nonlinear Science, 7(4):520-543, 1997.
13

---

Published as a conference paper at ICLR 2025
William Peebles and Saining Xie. Scalable diffusion models with transformers In Proceedings of the JEEE/CVF International Conference on Computer Vision, pp. 4195 4205, 2023.
Yifan Peng. Siddharth Dalmia, lan Lane, and Shinji Watanabe. Branchformer: Parallel mlp-attention architectures capture local and global context for speech recognition and understanding. In International Conference on Machine Learning, pp. 17627-17643. PMLR, 2022.
René Pinnau, Claudia Totzeck, Oliver Tse, and Stephan Martin, consensus- based model for global optimization and its mean-field limit. Mathematical Models and Methods in Applied Sciences, 27
(01):183-204, 2017.
Alec Radford. Improving language understanding by generative pre-trainitt pre-training. 2018.
Geo rey Zweig and Mukund Padmanabhan. Exact alpha-beta computation in logarithmic space with application to map word graph construction. 2000.
Michael F. Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Sinkformers: Transformers with doubly stochastic attention. In Proceedings of The 25th International Conference Arti- ficial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pp. 3515-3530, 28-30 Mar 2022.
Toby Shevlane, Sebastian Farquhar, Ben Garfinkel, Mary Phuong, Jess Whittlestone, Jade Leung, Daniel Kokotajlo, Nahema Marchal, Markus Anderljung, Noam Kolt, et al. Model evaluation for extreme risks. arXiv preprint arXiv:2305.15324 2023.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- fional Conference on Learning Representations, 2020.
Kenneth Stanley, David D' Ambrosio, and Jason Gauci. A hypercube- based encoding for evolving large-scale neural networks. Artificial life. 15:185-214 02 2009. doi: 10.1162/arl.2009.152 15202.
L. Storm, H. Linander, J. Bec, K. Gustavsson, and B. Mchlig. Finite-time Iyapunov exponents of deep neural networks, 2023.
Philipp Stumm and Andrea Walther. New algorithms for optimal online checkpointing. SIAM Journal on Scientific Computing, (2(2):86688442200 2010.
Anh Tong, Thanh Nguyen- Tang, Toan Tran, and Jaesik Choi. Learning fractional white noises in neural stochastic differential equations. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: //pp/////eeeeeeeeeeeeeeee net
Anh Tong, Thanh Nguyen- eguyen-Tang, Dongeun Lee, Toan M Tran, and Jaesik Choi. Sigformer: Signature transformers for deep hedging. In Proceedings of the Fourth ACM International Conference on Al in Finance, ICAIF '23, pp. 124-132, New York, NY, USA, 2023, Association for Computing Machinery. ISBN 9798400702402.
doi: 10. 10.114573644233336666411 URL /tttps://doi...
orq/10.11453 604237. 2...11445//60022333326446611
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint vrXiv:2302.13977 2023.
Belinda Tzen and Maxim Raginsky. Neural stochastic differential equations: Deep latent gaussian models in the diffusion limit, 2019, URL https: /
awwiiiiivvvi.. vivi wttg/abs/////ooooooro/o////.aeer 905. 09883.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NewrlPS, 2017.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Ryan Vogt, Maximilian Puelma Touzel, Eli Shlizerman, and Guillaume Lajoie. On Iyapunov expo- nents for mns: Understanding information propagation using dynamical systems tools. Frontiers in Applied Mathematics and Statisties, 8:8 1:818799, 2022.

---

Published as a conference paper at ICLR 2025
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov. and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference Machine Learning, pp. 35151-35174, PMLR, 2023.
Johannes Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, et al. Uncovering mesa-optimizztion algorithms in transformers. arXiv preprint arXiv:2309.05855 2023.
Sarah Wiegreffe and Yuval Pinter. Attention is not not explanation. In Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 11-20, 2019.
James Hardy Wilkinson. The algebraic eigenvalue problem, volume 662. Oxford Clarendon, 1965.
Alan Wolf, Jack B Swift, Harry L Swinney, and John A Vastano. Determining Iyapunov exponents from time series. Physica D: nonlinear phenomena, 16(3):85-3177, 1985.
Tianjun Zhang, Zhewei Yao, Amir Gholami, Kurt Keutzer, Joseph Gonzalez, George Biros, and Michael Mahoney. Anodev2: A coupled neural ode evolution framework, 2019.
Yaofeng Desmond Zhong, Tongtao Zhang, Amit Chakraborty, and Biswadip Dey. A neural ode interpretation of transformer layers, 2022.
15

---

Published as a conference paper at ICLR 2025
A DETAIL ON TIME-DEPENDENT WEIGHTS
In the main text, we define the time-dependent weights as:
(t) Proj (MLP(Simusoida((())
Each component will be explained further follows. The temporal embedding is concatenation of input time and sine and cosine of the projection of t higher dimensions.
Sinusoidal(t) Concat(t, sin(,,), cos(ur)).
The multilayer perceptron is two-layer neural network with SiLU activation. The projection layer a linear layer with a reshape function. In our experiment, we set the weight vector w E RI22 as
i 127
log (10*)
128
The hidden layer dimension of the MLP block is the embedding dimension demb.
Note that our models use time-dependent weights for layer normalization, with time embedding serving as input.
Figure 6: Visualization of the output of Sinusoidal function.
B FURTHER DETAILS ON LYAPUNOV EXPONENT
B.1 BRIEF REVIEW ON LYAPUNOV EXPONENT
In essence, the Lyapunov exponent of dynamical system quantifies the trajectory separation rate. For dynamical system described by =(t) = f(((t), consider two trajectories z(1) (t) and 2² (t). The Lyapunov exponent roughly captures the separation at time t relation to the initial separa- tion, expressed as [z(1) (t) - x (t) = exp(At)]
- (0))
In practice, obtaining the spectrum of Lyapunov involves computing the eigenvalues of matrix 21 Y()Y (t) at the limit where (t) satisfies
(()) J()Y((),, -
where (() represents the Jacobian matrix of f(x) with respect to sr(). This Jacobian matrix de- scribes how the tangent vectors associated with the given trajectory evolve. The numerical method to compute the Lyapunov exponent can be found in Wolf et al. (1985); Holzfuss & Parlitz (1991); Datseris (2018).
B.2 DERIVATION OF JACOBIAN TERM IN COMPUTING LYAPUNOV EXPONENT
Consider the case that want to measure the sensitivity of the i-th token to the j-th token. We need to extract the Jacobian Jij of the vector field in Equation equation! 7.
To make notation readable, rephrase our goal measure sensitivity of the in-th position to the out-th position. We simplify the multihead setting to one-head setting.
We also denote
Q(t)x;(), k; = =(()aa(() e; V(t)x.tt).
The first component the Jacobian Jin.out now is O x'r))) and is computed as
f A(m)))
Zqg
O (Xout F'(n))
fIrout x'rr))
[gg Day
0ka
Bus
(Kouts [((x))
16

---

Published as a conference paper at ICLR 2025
The individual terms in the above equation computed as
8
exp(1onnhi)
[Tom F(m))
= 0,
Oqn
0qia
Lout
==1
8
exp(domkk) ki)
exp(drunhia)
fooous [lli))
Is = Am--ooutes
Utin
0win
Lout
Louu
==1
exp(aautti)
exp kin) T
exp()oukss) exp kin
Okin
[Tonu Flmi)
Okin
Louut
Vj
Laut
You
Σ
Louut
Louu
Uj Tour
j=l
Amn- Min
Σ A A---outtyll
Amm- Pin
f (Dout mmp) Toun
expla, k.)
Here, we denote the attention
Ai-s
with L, = E. exp(e] ke).
Also, we have
0qm
Okin
dem (t).
Q(1,,
K(t,,
Orin
Drin
Drin
The second component in the Jacobian Jin,ou ( and can be computed easily.
Discussion The Jacobian Jin comprises two components: the attention information represented by Ain-out and the change between input and output represented by Vin - [ Tpj)). This ap- proach differs from heuristic methods that utilize attention matrices to explain transformers, which have been criticized for their limitations (Jain & Wallace, 2019). By incorporating both attention information and input-output changes, method provides more principled approach to under- standing the behavior of transformer models.
C EXPERIMENTS
Here, we provide details about data and experiment set up, e.g., hyperparameters, optimizers.
C.1 WIKITEXT DATASET
We utilized the Wikitext dataset, accessible via the Hugging Face dlwh/wikitext_100 detokenizeed.
dataset ID:
C.2 OPENWEB TEXT DATASET
Following the configuration from Dao et al. (2022), partitioned the entire Open WebText dataset into training and testing sets. The test set corresponds to 0.0005 of the entire dataset.
C.3 MODEL SPECIFICATION
The details of model specification of GPT-small, GPT-medium, GPT-large and Llama- IB are de- scribed in Table 3.
Specification Parameters
GPT-small
GPT-medium
GPT-large
Alama-1B
T17M
345M
762M
1213M
Layers Hidden Size
12
24
36
16
768
1024
1280
2048
Attention Heads
12
16
20
16
Table 3: Specifications for GPT-small, GPT-medium, GPT-large models, and Llama-1B BB.
17

---

Published as a conference paper at ICLR 2025
C.4 HYPERPARAMETHR SETTING
All models were trained using the Adam optimizer, along with a weight decay of 0.1 and dropout rate of 0.1. For GPT-small and GPT-medium models, we used Adam's 's default hyperparameters (3) = 0.9, $2 = 0.999). Based on our findings in Section C.5, we modified these parameters for GPT-large and Ilama-1B 1B models, setting 31 = 0.9, f32 0.95. The number of warm-up steps was set to 1% of the total training steps. Additionally, the cosine learning rate schedule is used. The ratio between the minimum learning rate and the base learning rate was fixed at 0.1.
We mainly use two A100 80GB GPUs for training our models. However, owing to hardware con- straints, there are variations the batch sizes used for training GPT-large, GPT-medium and GPT- small on the OPENWEBTEXT dataset. Specifically, the batch size for GPT-medium and GPT-large is set to 256, while the batch size for Gmllsmall is configured at 512. Mcanwhile, the batch size for training on the WIKITEXTI01 103 dataset is set at 256 for all models. For experiments with the Llama- IB configuration, training was conducted using eight NVIDIA H100 GPUs, each with 80GB of memory.
Backpropagation We employ the optimal online checkpointing technique (Stumm & Walther, 2010) rather than opting for the aptimize-then-disreee in Chen et al. (2018). The latter methods may introduce error accumulation, leading to instability during training. For a thorough discussion of the advantages and disadvantages of these techniques in backpropagation for neural ODEs, approach shares similarities with the techniques used in training deep neural networks (Chen et al., 2016; Gruslys al., 2016; rey Zweig & Padmanabhan, 2000).
approach or the adjoint method, as outlined
can refer to (Kidger, 2022, Chapter 5). Notably, this
C.5 ON ADAM'S B2 PARAMETER
Since DITTEOFORMER' 's weights are generated by time-dependent weight functions, its parameter- ization differs from GPT models, leading to distinct loss landscapes and optimizer behaviors. In practice, we found that setting the Adam optimizer parameter 83 0.95 improves optimization stability and convergence for DIFFEOFORMER compared to the standard value of 0.999.
40
40
(8, 0.95)
GPT (8, 0.95)
Val perplexity
35
Ours (82 0.999)
Val perplexity
35
GPT (82 0.999)
30
30
25
25
10000
20000
30000
40000
10000
20000
30000
40000
Step
Step
(a) DIFFERRRRMER
(b) GPT
Figure 7: The effect of modifying Adam's 32 parameter to 0.95. (a) DIFFEOFORMER shows sig-
nificant performance improvement with Bg 0.95 (from 15.89 to 15.43). In contrast, while GPT initially converges faster with B3 = 0.95, its final performance is inferior to that achieved with 32 = 0.999. Both models exhibit more stable and smoother training curves with decreased B2 val- ues.
When comparing DIFFEQFORMER with GPT-large on the OPENWEBTHE dataset, using A2 = 0.95 reduced our model's 's validation perplexity from 15.89 to 15.43, surpassing GPT baselines. In contrast, GPT models perform better with 82 = 0.999, as shown in Figure 7. We observe that lower 32 values generally produce smoother learning curves in both models.
We hypothesize that this behavior arises because the parameters of time-dependent weights (the MLP block in W(''s architecture) require rapid adaptation generate flexible weights. Since
18

---

Published as a conference paper at ICLR 2025
higher B2 values maintain longer memory of past squared gradients. lower 32 enables quicker adaptation of dinendepedennt weights by reducing the influence of historical gradients. This can explain the improvement for our model shown in Figure 7.
C.6 OTHER BASELINE ARCHITECTURES
We conduct additional experiments to evaluate the adaptabilityy of DIFFEOFORMER across differ- ent architectures, Specifically, we implement a LLaMA-based version of DIITEOFORMER follow- ing the architecture proposed by Touvron al. (2023). The key architectural differences between LLaMA- based and GPT-based implementations include: (1) the use of rotary positional embed- dings (RoPE) instead of absolute positional embeddings, (2) RMSNorm for layer normalization rather than the standard LayerNorm, and (3) SwiGLU activation functions in feed-forward blocks instead of ReLU. Importantly, these architectural modifications do not alter the core principles and effectiveness of DIFFEMFRRMRR' design.
20
Ours (Llama-1B)
Val perplexity
18-
Llama-1B
16-
14
12
10
8-
10000
20000
30000
40000
Step
Figure 8: Comparison of validation perplexity between our model and Llama- both trained with Adam optimizer (31 = 0.9, 82 = 0.95).
Figure 8 compares the validation perplexity between our Llama-based DIFFEOFORMER and the baseline Llama-IB IB on the OPENWEBTEXT dataset. Our model consistently achieves lower per- plexity, demonstrating DisTEOFORMER'S effectiveness when implemented with modern architec- tures like Llama.
C.7 DOWNSTREAM TASK EVALUATION
We evaluate the performance of DIFFEOFORMER and baselines in the setting of GPT-large and Llama-IB. We use the framework Im-evall uat siiiionnneeee the pre-trained models. The performance is assessed with the tasks having the similar setup with Bi- derman et al. (2023). Figure 9, Figure 10, and Table 4 show the performance compared against the GPT-large baseline actoss different tasks in two settings: zero-shot and five-shot. Figure 11, Fig- ure 12, and Table 5 show the performance compared against the Llama- 1B baseline across different tasks in two settings: zero-shot and five-shot.
(Gao et al., 2024) to evaluation
GPT-large Our model demonstrates competitive performance against GPT--arge across multiple benchmarks. In zero-shot evaluations (Figure 9), we achieve statistically significant improvements on reading comprehension tasks: Lambada OpenAl and Lambada Standard. Additionally, we ob- serve modest but consistent gains reasoning tasks including LogiQA, PIQA, and SciQ. In five- shot settings (Figure 10), our model outperforms GPT-large on 5 out significant improvements Lambada Standard and Winogrande. Performance remains comparable on two tasks, while GPT-large shows marginal, though not statistically significant, advantages on two others. These results suggest our neural ODE architecture maintains exceeds GPT-large's capabilities across diverse language understanding and reasoning tasks.
tasks, with statistically
Llama-1B
Our model demonstrates notable improvements over Llama- 1B in both zero-shot and five-shot settings, shown in Figures 11 and 12. The improvements particularly pronounced in
19

---

Published as a conference paper at ICLR 2025
reading comprehension tasks observe statistically significant accuracy gains (2% 4%) on both Lambada OpenAl and Lambada Standard benchmarks. These results suggest our neural ODE-based transformer architecture enhances language understanding capabilities while maintaining competi- tive performance across other evaluation metrics.
Model Performance Compartion
Performanc: Difference (Our Medel CPT)
20
Heintrrad
Lasis da Opetad
Lactadd srdddrr
Pose delp
ARc-Callep
Bay
Weinebeed
Lasila da Operad
Landaade Sradddd
Lauiti
8:g
AlC
Figure 9: Benchmark score for zero-shot in GPT-large configuration.
Modal Performence Comperioon
Oue Madel
10
Hellereed
Laaiiata Orenall
Lentanie Standerd
Learche
Pida
coo Intourned
ABE Chathoge
Healatvedd
Lambaad Openal
Lisst Staaderd
Liaetta
ANsEL
Figure 10: Benchmark score for five-shot GPT-large configuration.
Model Performance Compartion
Performance Difference (Dur Model Liama)
Helntrood
Lachasts Deread
Lasbaad Stanserr
pose
warrrr
#Rc-chaleus
La opecaed
Lassiaas steacdedd
Figure 11: Benchmark score for zero-shot in Llama- 1B configuration.
20

---

Published as a conference paper at ICLR 2025
Model Performance Comportoon
Performance Differonce (Dur Model -Lamm))
60-
10
Lecandla (peridd
Lantaade Standerd
HAMEO
Intepraaee wartrr
ARC mecccaatoon
Bay
Lana vol Opreed
Hollchied
Sausderd
Loutta AisEo
Pite
warcrry
ARc-
ABC
Figure 12: Benchmark score for five-shot in Llama-1B configuration.
Zero-shot
Five-shot
Benchmark
GPT
Our Model
GPT
Our Model
ARC-Challenge
20.65 (1.18)
20.90 (1.19)
21.84 (1.21)
21.25 (1.20)
ARC-Easy
36.03 (0.99)
35.82 (0.98)
$7.19
36.32 (0.99)
HellaSwag
25.53 (0.44)
25.61 (0.44)
(0.43)
25.52 (0.44)
Lambada OpenAl
37.24 (0.67)
38.39 (0.68)
27.79 (0.62)
27.83 (0.62)
Lambada Standard
29.30 (0.63)
30.37 (0.64)
23.05 (0.59)
24.82 (0.60)
LogiQA
21.35 (1.61)
21.35 (1.61)
20.74 (1.59)
21.81 (1.62)
MMLU
22.95 (0.35)
22.99 (0.35)
26.08 (0.37)
25.10 (0.36)
PIQA
56.26 (1.16)
56.96 (1.16)
55.82 (1.16)
56.69 (1.16)
SciQ
45.70 (1.58)
47.20 (1.58)
50.30 (1.58)
50.90 (1.58)
Winogrande
50.75 (1,41)
52.57 (1.40)
49.33 (1.41)
53.67 (1.40)
WSC273
55.31 (3.01)
53.11 (3.03)
50.92 (3.03)
53.11 (3.03)
Table 4: Performance comparison between GPT and our model in zero-shot and five-shot settings.
Zero-shot
Five-shot
Benchmark
Llama- 1B
Our Model
Llama-IB
Our Model
ARC-Challenge
22.27 (1.22)
22.27 (1.22)
21.76 (1.21)
22.44 (1.22)
ARC-Easy
29.67 (0.94)
31.10 (0.95)
30.01 (0.94)
31.44 (0.95)
HellaSwag
25.20 (0.43)
25.32 (0.43)
25.19 (0.43)
25.36 (0.43)
Lambada OpenAl
41.84 (0.69)
43.74(0.69)
30.47 (0.64)
34.41 (0.66)
Lambada Standard
27.58 (0.62)
31.55 (0.65)
24.04 (0.60)
28.22 (0.63)
LogiQA
21.51 (1.61)
22.27 (1.63)
21.04 (1.60)
21.20 (1.60)
MMLU
22.93 (0.35)
22.98 (0.35)
26.43 (0.37)
24.72 (0.36)
PIQA
53.26 (1.16)
53.81 (1.16)
53.65 (1.16)
53.10
SciQ
31.70 (1.47)
(1.49)
34.20 (1.50)
34.80 (1.51)
Winogrande
50.51 (1.41)
49.64 (1.41)
51.38 (1.40)
49.33 (1.41)
WSC273
53.11 (3.03)
52.38 (3.03)
52.01 (3.03)
50.92 (3.03)
Table 5: Performance comparison between Llama and our model zero-shot and five-shot settings.
C.8 ON TRAJECTORY SIMULATION
This simulation tries to replicate scenarios following the spectral dynamic of trained DIFFEQ- FORMER: increasing the magnitude of eigenvalues with peak near the last layer.
We consider the attention-only model with head,
21

---

Published as a conference paper at ICLR 2025
(.( (() F/ (0,)),,,
E [0. T},
2,(0) =X,;
with
(Q(t)z,(t), (t)) (Q));))))))()))))))))))
Σe exp
v())a(())
vd
j=1
The weights are randomly initialized at time 0 and time-dependent under the following dynamies:
Q(1) k(() Ato(() V(t) Vof((t)
A 1),,),, vod N(0,1).
where f()) controls the magnitude of the matrices over time. We test different configurations of ft) as polynomials including:
1 [o(1)
1 t 3 Sa(1) 2 T
2'
fi(t) 1
1
fa(t) = 2 T
2T
2
2
f2(1) 2 7
fa(() 2 T
In our simulation, we employed the following parameters: a time horizon T = 20 and an ODE solver step size 0.1, effectively resulting 200 layers of attention. We generated a sequence of length 40 with dimension 3, using uniform random sampling within the range [-2,
Figure 13 illustrates the trajectories of ODEs given the inputs. We observe that dynamics character- ized by functions increasing rates exceeding linear tend not to form clusters. Moreover, we noted that the output dispersion increases with the polynomial order the function. For instance, when fi is defined formation. In contrast, when fa is represented by fourth-order polynomial, its outputs demonstrate greater scatter.
linear function of t, the corresponding ODE trajectory exhibits incipient cluster
22

---

Published as a conference paper at ICLR 2025
0.52
0.48
10
[t) 0.5
0.2
(t) = 0.5t/20 20
0.2
(2(() = 0.5(1
0.2
(t( = 0.5(t/20)a
0.0
fa(t) = 0.5(t/20)4
0.2
0.8
15
ft(() = 0.5(1 - t/20)2
Figure 13: Trajectories of ODE dynamics for various settings of ft)).
23

---

Published as a conference paper at ICLR 2025
D ABLATION STUDY
Here we study our model when varying embedding dimension. Then we investigate an alter- native architecture for time-dependent weights.
D.1 ON THE DIMENSION OF TIME EMBEDDINGS
In equation 6, demb denotes the dimensionality of the time bedding used to derive the time-dependent weight (t). This hyperparameter directly influences the model's parameter count, with higher values potentially leading to increased model complex- ity. Our investigation aims assess the impact of varying demb on model performance. Figure 14 illustrates the validation per- plexity of DIFFEQFORMER across different demb values, bench- marked against the GPT-large model. The results demonstrate a positive correlation between increasing demh and improved perfor- mance. For models trained on the OPENWED TEXT dataset, rela- tively high demb values 144 are necessary to achieve optimal perfor- mance. Moreover, models with larger hidden dimensions require correspondingly higher demb values. We posit that demb is related to the concept of intrinsic dimension, described by Li al. (2018). This connection stems from the fact that in equation 6, the time- dependent weight (L)) results from projecting low-dimensional vector of size demh-
Vall perplexity
25-
GP1
25
20006 15
5000
10000
15000
Step
Figure 14: Perplexity of DIF- FËOFORMER with varying demb-
Figure 15 shows the training throughput between two models. As demb increases, a slight decrease in throughput observed. However, the magnitude of this decrease is notably small, suggesting that DIFTEOFORMER maintains competitive throughput even with nigher-dimensiones time embedding. Despite DIFFEOFORMER having significantly higher parameter count compared to the GPT model (see Table 6) due to its time-dependent weights, this has a minimal impact on the overall throughput. The additional computational cost is primarily incurred during the precomputation of the weights prior to the forward pass. The forward pass itself has the same computational cost as that of the GPT model. For the same reason, the GPU usage of DIFFEQFOR can be similar to that vanilla transformers. We are able train DIFFEOFRMMRR with medium-sizer architecture using two NVIDIA A100 GPUs, despite having 1.8 billion parameters (demb 144).
Inference time While learning the time- denendepnndent weights (W) in transformers requires both forward and backward passes during training, this process becomes unnecessary during inference. Since the training is complete, we can precompute W (t) for each layer and reuse the computed weights when needed.
Examples per second
GPT2
Ours
demb
Number of parameters
20
300M
91
48
650M
72
960M
96
1250M
120
1560M
89
144
1860M
88
Table 6: Parameter count of DIF- FEOFORMER with medium-sized architecture.
20
48
72
96
120
144
dembed
Figure 15: Throughput comparison, measured in examples per second, between GPT and DIFFEOFORMER for differ- ent values of the time embedding dimension (demb).
24

---

Published as a conference paper at ICLR 2025
GPU memory usage Figure 16 shows the peak GPU memory usage across different demb values when training DIFFEOFORMER with GPT-medium architecture at batch size 256. While increasing demb expands the model'ss parameter count, it does proportionally increase GPU memory usage for multiplicative factors. This is because the time-dependent weights (()) are computed once per forward pass generate the weights. Such generate weights then interact with transformer hidden states across batches.
160000
GPT
8.6%
Peak GPU Memary (MB)
150000
Ours
5.1%
140000
4.2%
130000
120000
16.4%
110000
10
50
100 demo
150
Figure 16: Peak memory of GPU for different settings of demb- Note that
the GPU memory values are obtained after setting JAX environment
variable
XLA PYTHON CLIENT ALLOCATOR=platfom..
D.2 ALTERNATIVE DESIGN FOR IEEEDEPENDENT WEIGHTS
Figure 18 shows the performance between two designs of time-dependent weights: (i) our model shares only Sinusoidal; (ii) alternative design shares both Sinusoidal and MLP blocks (see Fig- ure 17). We observe the our model outperforms the alternative due to having non-linearity before the projection block.
Proj
Q(H)
Proi
Q())
MLP
Proj
K((t)
Proj
K(I)
Temperal
MLP
Proj
V(t)
Tempora
-LP
Proj
V(K)
MLP
Proj
Proj
#FF(()
Proj
FF
Prol
Figure 17: Left: Our model's time dimendenen weights: sharing Sinusoidal component for all weights Q(t), (t), ((t), WFF. FFF Right: An alternative architecture shares both Sinusoidal and MLP block.
25

---

Published as a conference paper at ICLR 2025
GPT
Smuseddil
Val perplexity
30
25
20
10000
20000
30000
60000
Step
Figure 18
E SPECTRAL DYNAMICS
This section presents the spectral dynamics for each dataset and model setting, as described in the main text.
E. ATTENTION BLOCKS
E.1.1 OPENWEBTEXT DATA
Head 0
Head
Head
Read 3
Head 6
Hoad
Head
Head
ToalAA
Head
Head
Read
Head
Head
Head
Reallh)
(owll)))
Reaun
Head
Resll()
Roalin
Layer
Layez
Layer
Layer
Layor
Laver
Layez
Layez
(a) QK pair
(b) OV pair
Figure 19: The spectral dynamics of the Query- Key (QK) and Output-Vaa Value (OV) pairs in the DIF- FEOFORMER model trained on the OPENWERTEXT dataset using medium-sized architecture.
26

---

Published as a conference paper at ICLR 2025
Head
Brad
Heed
Head 3
Head
Head
Head
Head
Boal(A
Heed
Head
Head
Head
Head
Head
Head
Head
Roal(e)
Hoad
Head
Head
Hoad
Head
Head
RowlD))
Layar
Layer
Layer
(a) QK pair
(b) OV pair
Figure 20: The spectral dynamics of the Query-Key (QK) and Output- Value (OV) pairs in the DIF- FEQFORMER model trained on the OPENWEB TEXT dataset using small-sizer architecture.
E.1.2 WIKITEXTIO3 DATA
Head
Head
Head
Head 3
Head
Head
Read
Head
Roall)
Hoad
Head
Bead
Head
Head
Head
Head
Head
Beed
Head
Roslla)
Layer
Layer
Layor
Layer
Layez
Layer
(a) QK pair
(b) OV pair
Figure 21: The spectral dynamics of the Query-Key (QK) and Output- Value (OV) pairs in the DIF- FEOFORMER model trained on the WIKITEXTI03 dataset using medium-siza architecture.
27

---

Published as a conference paper at ICLR 2025
Head
Head
Head
Bead
Head 0
Head
Head
Hewd
Bowl(A)
Head
Head
Head
Bead
Head 4
Head
Head
Head
Roal(e)
Roali
Head
Head
Head
(oealA)
Roal/A)
Layer
Layer
Layer
(a) QK pair
(b) OV pair
Figure 22: The spectral dynamics of the Query-Key (QK) and Output- Value (OV) pairs in the DIF- FEQFORMER model trained on the WIKITEXTI 03 dataset using small-sized architecture.
E.1.3 VANILLA TRANSFORMER
Head 0
Head
Head 2
Head
Head
Head
Head 2
Head 3
Real(a)
Real(k)
Hebd 4+
Head 5
Head
Head 74
Head 4
Hebd
HTlad 6
Head 7
RealA)
Real())
Head
Head
Head
Real(A)
Real())
Held
Head 13
Head
Hebd
Head
Heedd 13
Head
Head
Real(h)
Real(1)
Layer
Layer
Layer
Layer
Layer
Layer (b) OV pair
Layer
Layer
(a) QK pair
Figure 23: The spectral dynamics of the Query- (QK) and Output- Value (OV) pairs in the vanilla transformer model trained the OPENWERS WETTEXT dataset using medium-sized architecture.
28

---

Published as a conference paper at ICLR 2025
Heed
Hoad
Road
Head
Meed
Hoad
Noad
Hoad
Liger
Lager
Layer
Lager
Lager
Lagur
(a) QK pair
(b) OV pair
Figure 24: The spectral dynamics of the Query-Key (QK) and Output- Output-Value Value (OV) pairs in the vanilla transformer model trained the OPEN W TEXTEXT dataset using small-sizer architecture.
29

---

Published as a conference paper at ICLR 2025
F LYAPUNOV SENSITIVITY EXAMPLES
More examples for Lyapunov sensitivity can be found in the following figures.
(A) BIGBANG is of those musical entities that transcends language. It's one of those rarc groups that both innovates and defines the direction genre lakes, Covering sound that includes hip hop, R&B and electronic dance, BIGBANG and
(B) Average of all heads
BIGBANG is one of those musical entities that transcends language. It's one of those rare groups that both innovates and defines the direction genre takes. Covering sound that includes hip hop, R&B and electromic dance, BIGBANG and
Head 0
BIGBANG is of those musical entities that transcends language. It's of those rure groups that both innovates and defines the direction genre takes, Covering sound that includes hip hop, R&B and electronic dance, BIGBANG and
Head 1
BIGBANG is one of those musical entities that transcends language. It's one of those rare groups that both innovates and defines the direction genre takes. Covering sound that includes hip hop, R&B and electronic dance, BIGBANG and
Head 2
BIGBANG is of those musical entities that transcends language. It's of those nure groups that both innovates and defines the direction genre takes. Covering sound that includes hip hop, R&B and electronic dance, BIGBANG and
Head 3
BIGBANG is one of those musical entities that transcends language. It's one of those rare groups that both innovates and defines direction genre takes. Covering sound includes hip hop, electromic dance, BIGBANG
Head 4
BIGBANG is of those musical entities that transcends language. It's of those rare groups that both innovates and defines the direction genre takes. Covering sound that includes hip hop, R&B und electromic dance, BIGBANG and
Figure 25; The next word is its which is the possessive pronoun for BIGBANG, (A) This shows that the obtained Lyapunov sensitivity can identify the main subject for possessive pronoun. (B) Analysis of token relationships through attention matrices. While individual attention heads capture specific aspects of token dependencies, the aggregated attention matrix across all heads does not fully explain next-token predictions.
30

---

Published as a conference paper at ICLR 2025
(A) North Korca has detained another U.S. citizen, Korean American professor, bringing three the number of Americans being held in Pyongyang. The Swedish Embassy in Pyongyang, which represents U.S. interests there becausc the United States does not have diplomatic relations with
(B) Average of all heads
North Korea has detained another U.S. citizen, Korean American professor, bringing three the number of Americans being held Pyongyang. The Swedish Embassy Pyongyang, which represents U.S. interests there because the United States does not have diplomatic relutions with
Head
North Korea has detained another U.S. citizen, Korean American professor, bringing to three the number of Americans being held in Pyongyang. The Swedish Embassy in Pyongyang, which represents U.S. interests there because the United States does not have diplomatic relations with
Head
North Korea has detained another U.S. citizen, Korean American professor, bringing to three the number of Americans being held Pyongyang. The Swedish Embassy Pyongyang, which represents U.S. interests there because United States docs not have diplomatic relations with
Head 2
North Korea has detained another U.S. citizen, Korean American professor, bringing three the number of Americans being held Pyongyang. The Swedish Embassy in Pyongyang, which represents U.S. interests there hecause the United States docs not have diplomatic relations with
Head
North Korea has detained another U.S. citizen, Karean American professor, bringing to three the number of Americans being held in Pyongyang. The Swedish Embassy in Pyongyang, w which represents U.S. interests there because the United States does not have diplomatic relations with
Figure 26: The next word is North in North Korea. (A) We also can see that Lyapunov sensitivity
highlights the related words including "North", "Korean", "diplomatic" in this context. (B) Analysis of token relationships through attention matrices. We observe the same behaviors as the previous example.
31

---

Published as a conference paper at ICLR 2025
(A)
Sleep is present and tightly regulated in every vertebrate species which it has been carefully investigated, but what sleep is for remains mystery. Sleep also present in invertebrates, and extensive analysis in Drosophila melanogaster has shown that sleep in fruit flies shows most of the fundamental features that characterize sleep in mammals. In Drosophila, sleep consists of sustained periods of quiescence associated with an increased arousal threshold. Fly sleep is modulated by several of the same stimulants and hypnotics that affect mammalian
Average of all heads
(B)
Sleep is present and tightly regulated in every vertebrate species in which it has been carefully investigated, but what sleep is for remains mystery. Sleep also present in invertebrates, and an extensive analysis in Drosophila melanogaster has shown that sleep in fruit flies shows most of the fundamental features that churacterize sleep in mammals, In Drosophila, sloep consists of sustained periods of quiescence associated with an increased arousal threshold, Fly sleep is modulated by several of the same stimulants and hypnotics that affect mammalian
Head
Sleep is present and tightly regulated in every vertebrate species in which it has been carefully investigated, but what sleep is for remains mystery. Sleep also present in invertebrates, and an exlensive analysis in Drosophila melanogaster has shown that sleep in fruit flies shows most of the fundamental features that characterize sleep in mammals, In Drosophila, sleep consists of sustained periods of quiescence associated with increased arousal threshold, Fly sleep modulated by several of the same stimulants and hypnotics that affect mammalian
Head
Sleep present and tightly regulated in every vertehrate species in which it has been carefully investigated, but what sleep for remains mystery. Sleep also present invertebrates, and extensive analysis in Drosophila melanogaster has shown that sleep fruit flies shows most of the fundamental features that characterize sleep in mammals. In Drosophila, sleep consists of sustained periods of quiescence associated with increased arousal threshold. Fly sleep modulated by several of the same stimulants and hypnotics that affect mammalian
Head 2
Sleep present tightly regulated every vertebrate species in which has been carefully investigated, but what sleep remains mystery. Sleep also present invertebrates, and extensive analysis in Drosophila melanogaster shown that sleep fruit flies shows most the fundamental features characterize sleep in mammals. Drosophila, sleep consists sustained periods quiescence associated with increased arousal threshold. Fly sleep modulated by several the same stimulants and hypnotics that affect mammalian
The next word is sleep. (A) The Lyapunov sensitivity highlight the word "sleep" which is the main topic of the paragraph and appears at the beginning of the sentence. In this case, it may overestimate the sensitivity of the word "mammalian". (B) Analysis of token relationships through attention matrices. We observe the same behaviors as the previous example.
Figure 27:
32

---

Published as a conference paper at ICLR 2025
(A)
SEOUL (Reuters) South Korean Lee Sedol won his first match against computer program developed by a Google subsidiary Sunday in the ancient board game Go, denying a clean sweep for the artificial intelligence in rive-matchl series, The world's top Go player Lee Sedol reviews the match after the fourth
Average of all heads
SEOUL (Reuters) South Korean Lee Sedol his first match against computer program developed Google subsidiary on Sunday in the ancient board game Go, denying clean sweep for the artificial intelligence in five- match series, The world's top Go player Lee Sedol reviews the match after the fourth
Head
SEOUL (Reuters) South Korean Sedol his first match against computer program developed a Google subsidiary Sunday in the ancient board game Go, denying clean sweep for the artificial intelligence in five- match series, The world's Go player Lee Sodol reviews the march after the fourth
Head
SEOUL (Reuters) South Korean Sedol won his match against computer program developed a Google subsidiary on Sunday in the ancient board game Go, denying clean sweep for the artificial intelligence five- match series, The world's lop player Lee Sedol reviews the match after the fourth
Head
SEOUL SEOUI_(Reuters) South Korean Lee Sedol won his first match against computer program developed by a Google subsidiary on Sunday in ancient board game Go, denying a cleun sweep for artificial intelligence five- match scries. The world's top player Lee Sedol reviews the match after the fourth
Head
SEOUL (Reuters) South Koreun Lee Sedol won his first match against computer program developed by a Google subsidiary on Sunday in the ancient board game Go, denying a clean sweep for the artificial intelligence five- match series. The world's top player Lee Sedol reviews the match after the fourth
Head
SEOUL (Reuters) South Korean Lee Sedol won his first match against computer program developed by a Google subsidiary on Sunday in the ancient board game Go, denying a cleun sweop for the artificial intelligence in five- match series, The world's top Go player Lee Sedol reviews the match after the fourth
The next word is match. (A) This is interesting case where the output of Lyapunov sensitivity highlights the word "match" "five-match" and the word "after" which indicates order of the matches. (B) Analysis of token relationships through attention matrices. We observe the same behaviors as the previous example.
Figure 28:
33

---

Published as a conference paper at ICLR 2025
(A)
An ancient Roman merchant vessel has been discovered off Italian coastline, reportedly in such good condition that much of the food it carrying might still be intact in its storage jars. "There are some broken jars around the wreek, but we believe that most of the umphorac inside the ship still sealed and food-filled,' Lt. Col. Francesco Schilardi of the police divers' group told the BBC of the containers, Local fisherman first became aware of the wreck when pieces of pottery began turning up in their
Average of all heads
(B)
An ancient Roman merchant vessel has been discovered off the Italian coastline, reportedly in such good condition that much of the food was carrying might still be intact storage jars. There are some broken jars around the wreck, but we believe that most the amphorae inside the ship are still sealed and food-filled,' Lt. Col. Francesco Schilardi of the police divers' group told the BBC of the containers. Local fisherman first became of wreck when pieces of pottery began turning up their
Head 0
An ancient Roman merchant vessel has been discovered off the Italian coastline, reportedly in such good condition that much of the food it was carrying might still be intact in its storage jars. There some broken jars around the wreck, but believe that most of the amphorae inside the ship are still sealed and food-filled," Lt. Col. Francesco Schilardi the police divers' group told the BBC of the containers. Local fisherman first became aware of the wreck when pieces of pottery began turning up in their
Head 1
An ancient Roman merchant vessel has been discovered off the Italian coastline, reportedly in such good condition that much of the food was carrying might still be intact in storage jars. There are some broken jars around the wreck, but we believe that most the amphorae inside the ship are still scaled and food-filled," Lt. Col. Francesco Schilardi of the police divers' group told the BBC of the containers. Local fisherman first became aware of wreck when pieces of pottery began turning up their
Head 2
An ancient Roman merchant vessel has been discovered the Italian coastline, reportedly in such good condition that much of the food it was carrying might still be intact in its storage jars. There some broken jars around the wreck, but believe that most of the amphorae inside the ship are still scaled and fild-filed,'' police divers' group told the BBC of the containers. Local fisherman first became aware of the wreck when pieces of pottery began turning up in their
Col. Francesco Schilardi the
Figure 29; The next word is nets. (A) Lyapunov sensitivity gives a high value for the word "in" which preposition and relevant to the next work "nets". Another relevant word "fisherman" is highlighted together with the object "pottery" in the nets. (B) Analysis of token relationships through attention matrices. We observe the same behaviors as the previous example.
34


Reviews:
[{"review_id": "4WSRHnEjSt", "reviewer": "Reviewer_Pfg2", "summary": "This paper proposes a novel approach to modeling transformer architectures using depth-varying neural ordinary differential equations (ODEs). The authors introduce time-dependent weights for attention and feed-forward blocks, implemented through neural networks that take a continuous layer index ($t$) as input. They analyze the model's dynamics through spectral analysis and Lyapunov exponents, aiming to provide insights into transformer behavior. The work also explores adaptive fine-tuning capabilities enabled by the ODE formulation, showing comparable performance to GPT models across various configurations and datasets.", "questions": "1. Could you provide theoretical justification for the specific design of the time-dependent weight parametrization? What motivated these particular choices, and how do they relate to the model's performance?\n2. The paper presents both analytical tools (spectral analysis, Lyapunov exponents) and empirical results, but the connection between them isn't clear. How do the insights from your analysis inform or explain the observed performance in the experiments?\n3. The adaptive fine-tuning capability seems promising, but its benefits aren't fully explored. Could you elaborate on scenarios where this flexibility provides concrete advantages over traditional fine-tuning approaches?\n4. Given that the Lyapunov exponent analysis is input-dependent, how generalizable are the insights gained from this analysis? What patterns or principles have you observed across different classes inputs that might provide more universal understanding of transformer behavior?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper makes an interesting attempt to bridge the gap between theoretical analysis of transformers and practical implementations by using neural ODEs with depth-varying weights.\n- The proposed adaptive fine-tuning approach, which leverages the discussed ODE formulation to enable flexible architecture changes during fine-tuning, presents a potentially useful direction for model adaptation.\n- The presentation includes detailed visualizations of spectral dynamics and Lyapunov exponent analysis, which helps readers understand the model's behavior.", "weaknesses": "1. Limited Technical Novelty:\n- The core formulation of transformers as ODEs is directly borrowed from Lu et al. (2019) without significant theoretical advancement.\n- The analysis methodology for the dynamics closely follows Geshkovski et al. (2023a,b) without introducing substantial new analytical tools or insights.\n- The time-dependent weight parametrization, while interesting, lacks theoretical justification for its specific design choices.\n2. Insufficient Analysis Rigor:\n- The spectral analysis lacks formal theoretical guarantees or concrete connections to model performance.\n- The Lyapunov exponent analysis is essentially equivalent to analyzing the Jacobian of transformer blocks, which varies with each input and doesn't provide generalizable insights.\n- Many claims about model behavior are supported by cherry-picked empirical observations without theoretical justification.\n3. Experimental Limitations:\n- The performance improvements over baseline models are modest and inconsistent across different settings.\n- The ablation studies don't sufficiently justify the architectural choices, particularly the dimension of time embeddings.\n- The paper oscillates between being an analysis paper and an empirical paper without excelling in either direction.\n\n4. Presentation Issues:\n- The paper structure is somewhat disjointed, with early sections (1-4) largely reiterating existing work.\n- Many design choices appear arbitrary and lack proper motivation or ablation studies.\n- The relationship between the theoretical analysis and practical benefits isn't clearly established.", "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": "EUyquclm7i", "reviewer": "Reviewer_4Us1", "summary": "This paper presents a neural ODE-based approach for generating transformer weights. Unlike existing methods that use weight-sharing, this approach employs hyper-networks to produce the attention and feed-forward layer weights as functions of a continuous layer index (or time).", "questions": "Could the authors include evaluation on downstream tasks? This information would be valuable in assessing how well the model generalizes beyond the initial validation perplexity. They can show some zero and 5 shot evaluation on benchmarks from appendix of [2].\n\n[2] Biderman, Stella, et al. \"Pythia: A suite for analyzing large language models across training and scaling.\" International Conference on Machine Learning. PMLR, 2023.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) This paper employs flexible computation and dynamically adjusts the step size of its ODE solvers to fine-tune the model with variable layer lengths, achieving competitive performance with vanilla transformers. This adaptation is particularly promising as it allows for post-pretraining adjustments to the model's architecture, which is a great direction for future research.\n\n2) They demonstrate the model's capacity through spectral analysis, revealing that increasing eigenvalue magnitudes inhibit cluster formation compared to weight sharing based prior works. \n\n3) They also leverage the Lyapunov exponent to analyze token-level sensitivity for explainablity where they utilises both attention map and Jacobian which has similarity with the work in computer vision[1] for better faithfullness.\n\n[1] Chefer, Hila, Shir Gur, and Lior Wolf. \"Transformer interpretability beyond attention visualization.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.", "weaknesses": "1) No evaluation of downstream performance is conducted.\n\n2) Training is computationally and memory-intensive; for instance, training a GPT-medium model requires 1.8B (with d_emb=144) parameters to achieve competitive performance, while the original GPT-2 medium model has only 345M parameters.", "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": "jqmB0lgBSM", "reviewer": "Reviewer_rn5p", "summary": "This paper presents DIFFEQFORMER, an innovative framework that reformulates transformer architectures through the lens of non-autonomous neural ordinary differential equations (ODEs). The approach conceptualizes transformer layers as a continuous-time dynamical system, with the distinguishing feature of parameterizing attention and feed-forward weights as continuous functions of layer depth via neural networks.The authors conduct rigorous spectral analysis of the attention mechanisms, demonstrating increasing eigenvalue magnitudes across layers—a finding that challenges current theoretical frameworks predicated on weight-sharing assumptions. They further advance model interpretability by incorporating Lyapunov exponent analysis for examining token-level sensitivity.", "questions": "The time-dependent weight parameterization in your approach introduces additional parameters compared to standard transformer architectures. While you demonstrate comparable throughput, I'm concerned about the memory efficiency. Could you clarify the actual memory overhead during training and inference? Specifically, how does the memory usage scale with model size and demb.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper is exceptionally well-organized with a clear progression from theory to implementation. Complex mathematical concepts are presented in an accessible manner with effective visualizations.\n2. The paper introduces an innovative approach using non-autonomous neural ODEs with time-dependent weights to model transformers, effectively overcoming limitations of previous weight-sharing methods. The comprehensive analysis through spectral dynamics of QK/OV pairs and investigation of clustering behavior provides deep theoretical insights into transformer dynamics while achieving competitive performance.\n3. The introduction of Lyapunov exponents for analyzing token-level sensitivity is highly innovative. This provides new tools for model interpretability that go beyond traditional attention visualization methods.", "weaknesses": "1. The model's performance degradation at GPT-large scale raises concerns about its viability for larger architectures. The lack of experiments on billion-parameter scale models leaves uncertainty about whether this approach can effectively scale to sizes relevant for state-of-the-art language models.\n2. While perplexity results are promising, the paper lacks zero-shot evaluation results which are crucial for understanding the model's practical capabilities. Additional metrics would better demonstrate the framework's advantages over traditional transformers.\n3. The performance appears sensitive to the $d_{emb}$, but the paper provides limited guidance on how to optimally set this hyperparameter for different model scales.", "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": "wy8cBg239Z", "reviewer": "Reviewer_hGru", "summary": "This work proposes an architecture that implements a Transformer using a Neural ODE approach (DiffEqFormer). The main contribution over existing Neural ODE works is that DiffEqFormer has time-dependent weights, as opposed to traditional weight-sharing. This equips DiffEqFormer with the interesting property of adaptive fine-tuning, i.e. baing able to fine-tune a model with N layers as a M layer version of itself. The authors show competitive performance compared to GPT-S/M/L on two datasets. Additionally, interpretability results are obtained using the Lyapunov exponent of the proposed dynamical system.", "questions": "**Questions**:\n\n* I encourage the authors to add a section (main body or appendix) about the computational differences betwee DiffEqFormer and a vanilla Transformer. For example, ML practitioners could be interested in training time, total number of parameters, FLOPs, GPU memory required and any other significant parameter the authors might think of, given their expertise in the field.\n\n* In Section 5.2, I suggest including interpretability results using attention matrices, to make the differences claimed evident to the reader.\n\n* In Section 6, can the authors elaborate on the choice of experimental setup? Why is GPT2 chosen among the breadth of language models available? Unless the constraints are compute-related, I believe this work could be much more impactful if DiffEqFormer was compared to a recent LLM architecture.\n  * If some recent LLM was used, I am really intrigued in knowing about other metrics related to \"_emergent capabilities_\" are for DiffEqFormer. For example, does it perform better at 0-shot reasoning, on MMLU, etc? I am not asking the authors to provide all this analysis, but I definitely encourage them to include at least a comparison with a recent LLM.\n\n* As far as I understand, the model is trained at fixed time intervals. I wonder whether the authors have thought about training at randomized time steps, as it is typically done for diffusion models.  Is it even feasible? Could this allow a more flexible adaptive fine-tuning later on? I would appreciate some discussion on this aspect during the rebuttal, for my better understanding of the work.\n\n**Comments**:\n\n* Overall, I believe this is a nice piece of research work. I enjoyed the reading and learnt throughout. My main concern is about the experimental setup. A comparison with some modern LLM would be very important to convince the community that DiffEqFormer is _equivalent_ to Transformer with the additional unique advantage of adaptability+interpretability. I am completely open to increase my review score upon a fruitful rebuttal.\n\n* I would like to share my opinion about the experiments in Tables 1-2. Although the authors try to explain why GPT2-large is slightly better in terms of PPL for OpenWebText, I believe the strength of these results lie in DiffEqFormer being similar (not necessarily better) to GPT. \n\n**Suggestions for readability + typos:**\n\n* Very minor: Fig 1-b could have the arrow sin blue, as a connection with what is shown in Fig 1-a.\n* L304: The OV matrix ~has~ have eigenvalues\n* L322: model ~collapse~ collapses\n* Fig 6 (legend) ~cratch~ scratch", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "**Originality:**\n\n* The time-dependent implementation proposed is both elegant and original in my opinion. \n* Using the Lyapunov exponent for interpretability is an original idea, which is also very well adapted to the framework proposed.\n* The idea of adaptive fine-tuning is also original to me, specially given how natural it is under the DiffEqFormer formulation.\n\n**Quality:**\n\n* The experimental setup related to interpretability of attention matrices and eigen-values is very complete. \n\n**Clarity:**\n\n* The paper is well written, with clear language. The mathematical notation and formulation is also clear and easy to follow. I want to congratulate the authors for that, I honestly was hesitant before reading the paper since many ODE papers are too heavy in notation. I really enjoyed reading this manuscript!\n\n**Significance:**\n\n* I believe this Neural ODE approach can have impact and can serve as starting point for further research. Both the idea of adaptive fine-tuning and time-dependent weights have strong research potential.", "weaknesses": "**Quality:**\n\n* The experimental setup related to Language Modelling (6.1) uses arguably old models (GPT2). Although the experiments are useful to understand how DiffEqFormer is able to match performance, I missed a comparison with a current SOTA model. I believe this work could have more impact if a recent model was included, not needing an extremely large model, but maybe a 1B parameter LLM. For example,  [Llama-3.2-1b](https://huggingface.co/meta-llama/Llama-3.2-1B) or [Gemma-2-2b](https://huggingface.co/google/gemma-2-2b) could be excellent candidates.\n\n**Clarity:**\n\n* I felt that the discussion about _attractive_ and _repulsive_ scenarios (5.1) is a little bit disconnected from the rest, and much harder to parse. For example, I could not directly understand why a $V=-I_d$ behavior leads to model collapse. Also, in the conclusion of this section (L336) the authors emphasize that \"_positive values from the QK pairs in Figure 2a and the copying behavior observed in the OV pair_\" hint to the presence of induction heads in the last layer(s). However, I understood that what is important is the magnitude of QK pairs, and the positiveness of OV pairs. There is a mismatch between this conclusion and the claims in the explanation.  I encourage the authors to revisit this section and provide a clearer explanation of why the QK, and specially the OV analysis, is important.\n\n**Significance:**\n\n* As said, this work can be impactful given the applicability of DiffEqFormer. However, lacking a comparison with SOTA LLMs can diminish such impact.", "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)