Disappearance of Timestep Embedding: A Case Study on Neural ODE and Diffusion Models
Authors: Bum Jun Kim
Yoshinobu Kawahara
Sang Woo Kim
Venue: Accepted by TMLR ()

Abstract: Dynamical systems are often time-varying, whose modeling requires a function that evolves with respect to time. Recent studies such as the neural ordinary differential equation proposed a time-dependent neural network, which provides a neural network varying with respect to time. However, we claim that the architectural choice to build a time-dependent neural network significantly affects its time-awareness but still lacks sufficient validation in its current states. In this study, we conduct an in-depth analysis of the architecture of neural ordinary differential equations. Here, we report a vulnerability of vanishing timestep embedding, which disables the time-awareness of a time-dependent neural network. Specifically, we find that the ConcatConv operation, which is widely used in neural ordinary differential equations, causes an additive effect of timestep embedding, which is readily canceled out by the subsequent batch normalization. This vanishing timestep embedding also arises for group normalization and is analyzed thoroughly with respect to the number of channels, groups, and relative variance. Furthermore, we find that this vulnerability can also be observed in diffusion models because they employ a similar architecture that incorporates timestep embedding to discriminate between different timesteps during a diffusion process. Our analysis provides a detailed description of this phenomenon as well as several solutions to address the root cause. Through experiments on neural ordinary differential equations and diffusion models, we observed that ensuring alive time-awareness via proposed solutions boosted their performance, such as classification accuracy, FID, and inception score, which implies that their current implementations lack sufficient time-dependency.

Published in Transactions on Machine Learning Research (05/2025)
Disappearance of Timestep Embedding: A Case Study on Neural ODE and Diffusion Models
Bum Jun Kim
bannjn.kkin@uuchbbb.....ooooo
Graduate School of Engineering, The University of Tokyo
Yoshinobu Kawahara
Graduate School of Information Science and Technology, The University of Osaka RIKEN Center for Advanced Intelligence Project
Sang Woo Kim
swkim eukim@postech edu
Department of Electrical Engineering, Pohang University of Science and Technology
Reviewed OpenReview: https: o/pereeveeen net/ per/forum? id=bpaLYaafog
Abstract
Dynamical systems often time-narying, whose modeling requires function that evolves with respect to time. Recent studies such as the neural ordinary differential equation pro- posed a time-dependent neural network, which provides neural network varying with respect to time. However, we claim that the architectural choice to build a time-dependent neural network significantly affects its time-awaree -waaenesss but still lacks sufficient validation in its current states. In this study, we conduct an in-depth analysis of the architocture of neural ordinary differential equations. Here, we report a vulnerability of vanishing timestep embed- ding, which disables the time-awareness of time-dependent neural network. Specifically, we find that the ConcatConv operation, which is widely used in neural ordinary differential equations, causes an additive effect of timestep embedding, which is readily canceled out by the subsequent batch normalization. This vanishing timestep embedding also arises for group normalization and is analyzed thoroughly with respect to the number of channels, groups, and relative variance. Furthermore, we find that this vulnerability can also be observed in diffusion models because they employ similar architecture that incorporates timestep embedding to discriminate between different timesteps during diffusion process. Our analysis provides detailed description of this phenomenon well as several solutions neural ordinary differential equations
to address the root cause. Through experiments and diffusion models, we observed that ensuring alive time-awareness via proposed solutions boosted thoir performance, such as classification accuracy, FID, and inception score, which implies that their current. implementations lack sufficient Lime-dependency.
1 Introduction
Deep neural networks have demonstrated remarkable potential in a wide variety of fields, including com- puter vision and natural language processing. Furthermore, recent studies have extended their usage to the modeling of time-varying dynamical systems. A prime example is the neural ordinary differential equation (NODE) (Chen et al.] 2018), which employs time-dependent neural network to model the differential equa- tion of a time-varying dynamical system. The NODE is able to describe complex dynamical system via a single neural network that varies with respect to time and has been deployed in several fields (Dupont et al. 2019
Yan et al. 2020 Norcliffe et al. 2020 Kidger et al. 2021 Xia et al. 2021
Another famous use of time-dependent neural networks in diffusion models (oo al. 2020 Song et al.
2021ct. Recent studies generative models cast. their task diffusion process, whose modeling requires

---

Published in Transactions on Machine Learning Research (05/2025)
discrimination between different timesteps throughout the process. For this task, time-dependent neural network has been employed, which is suitable to describe different phenomena evolving with respect to time.
Despite different contexts for the emergence of NODE and diffusion models, we discover that they share a similar design pattern in architecture: The time-dependent neural network incorporates time-awareness by introducing a module of timestep embedding, whose values vary with respect to time. From this observation, we collectively analyze the two families of models in terms of their architectural configurations, which are subjects of this study.
Albeit. the wide usage of NODE and diffusion models for their application, however, we find that no study has yet been conducted to investigate the validity of the current architectural choice. In other words, we claim that the architectural choice to build a time-dependent neural network still lacks sulficient validation. The lack of architectural analysis of these time-dependent neural networks may cause significant problems: For example, if the architecture has a certain vulncrability, it would perform unwanted bchavior and provide degraded performance.
This study inspects the time-dependency of NODE and diffusion models. Eventually, we report a vulner- ability in their architecture: The timestep component is inherently prone to vanish, which disables time- awareness. This phenomenon was observed for both NODE and diffusion models owing to their similar design patterns. We provide a unified viewpoint on the vanishing timestep embedding of NODE and diffusion mod- els, along with a detailed explanation of their architecture. We also present several solutions that address the root cause of this problem to guarantee sufficient time-dependency. Through practical experiments on NODE and diffusion models, we observed that encouraging a sufficient timestep embedding significantly boosted their performance.
2
Vanishing Timestep Embedding
2.1 Background: How NODE incorporates time-awareness
The time-dependent neural networks we discuss indicate a family of neural networks that receive both input feature and timestep. In contrast. to the common neural networks that provide fixed output without time- dependency, time-dependent neural network provides a time-varying function, which suitable for modeling different phenomena that vary with respect to time. Their prime examples in the current rescarch community include NODE (Chen et al. 2018) and diffusion models (Ho el al., 2020 which are the main subjects of our study. In this section, we first review NODE, state its problem, and then extend our claim to diffusion models.
The inspiration behind NODE is to understand the depth of neural network as time and the residual network (He et al., 2016) as a discrete-time function. In NODE interpretation, a feature map evolves with respect to time in a dynamical system to be an output. They propose continuous- time neural network that is formulated in terms of an ordinary differential equation (ODE)
dh(t) = ((b(),,,0)
(1)
dl
While NODE f is a neural network parameterized with 0, it receives both input feature h(t) and timestep E T), providing a different output depending on the timestep t. Applying a suitable ODE solver to f yields an integration h(T)) from the corresponding initial value h(0). Note that the term "time" in our study differs from the conventional notion of time used in time-series data. In NODE studies, the time unit is something similar to a depth unit in neural network, not indicating the real notion of time, while time-scries data can be used as input for NODE. The behavior of time-dependent function evolving in I E [o can be understood as a behavior of data passing through the layers of a neural network, from the first layer to the last layer, e.g., the 101st layer. In other words, the behavior of a NODE block being time-dependent is similar to a residual network having different functions or weights across different layers. Thus, repeating a NODE block with respect to different times similar to passing a residual network across whole layers, which enables performing a static task such image classification and generation. Tf NODE
2

---

Published in Transactions on Machine Learning Research (05/2025)
Table 1: List of notations used in this study.
Notation
Meaning
f
NODE that is modeled as a neural network.
Paramelers of NODE.
h(t)
Input feature NODE. Timestep in NODE. Ranges from to T. An imput feature map NODE pipeline.
X
M.W.C
Height, width, and the number of channels for the input feature map X. All-ones tensor. All elements one. Output of ConcatConv in NODE pipeline. The k-th channel of Z.
Z
z²
W
Convolutional kernel used in ConcatConv. Convolutional operation.
[X:J]
Concatenation of X and tJ with respect the channel dimension. Subset of W from the first the C-th channel. Subset of W corresponding (C + 1)-th channel. Output of convolution an all-ones tensor with the kernel Wc+1. Timestep embedding.
W WI.c
W C+1
V
v J
An element of v that corresponds to the (i,j) coordinate and the kth channel. Mean and variance computed through BN manner.
EHJ Varl
Pr
Positional timestep embedding.
By
Partial derivative of with respect to æ.
Dy x Dw
The spatial size of the kernel. Variance of weight initialization.
LN
Layer Normalization. Group Normalization.
GN
does not have time dependency, it becomes fixed-weight RNN with residual connection. This behavior is further analyzed later.
Here, we describe the NODE for vision tasks. For an image classification task, an input image is processed by an early-stage such as strided convolution to produce an image feature, which acts as the initial value for the ODE, f.c., h(0). Applying NODE yiclds h(T), which corresponds to the output of the continuous- depth neural network. A head, such as a fully connected (FC) layer and softmax, outputs a final score of classification probability, In other words, adopting NODE means replacing a sequence of residual blocks with single ODE block while using the remaining layers such the carly-stage and head.
To build a time-dependent neural network f, the original NODE study proposed pipeline of [GN-AC Act- ConcatConv GN- Act ConcatConv GNJ. The ConcatConv operation also referred to as Conv2dTime is key to incorporating lime-dependency. Consider an input feature map X E H XC where H, W, and C correspond to the spatial size and the number of channels of the feature map. The ConcatCony operation generates a timestep feature iJ E H where nated to the input feature map with respect to the channel dimension, yielding [X; (J] E H Now, vanilla convolutional operation is applied to produce Z W [X: tJ] II #xxXC where indicates the convolutional operation and W indicates the convolutional kernel. The convolution here blends the timestep feature into all channels. With rich nonlincarities such as group normalization (GN) (Wu & Hc 2018 and ReLU activation function (Act), the output becomes lime-dependent. Indeed, the pipeline with Concal.Covv is considered the most standard practice in Neural ODE. Since its introduction, this pipeline along with ConcatConv has been widely adopted in follow-up studies on NODE (Dupont et al. 2019 Yan et al. 2020; Norcliffe et al. 22020] Kidger et al., 2021; Xia al. (022)). Although these studies do not explicitly state adopt the pipeline with ConcatCCov, we find that is implicitly used in their source code.
an all-ones tensor. The timestep feature concate-
3

---

Published in Transactions on Machine Learning Research (05/2025)
*
X
w 1.C
z
*
t]
W
3
-3-2-1
33
t
*
2
t
2
w
Figure 1: In Concat Conv operation, applying convolutional kernel W C+1 to IJ equivalent to using lv that has the same element spatially
2.2 Problem statement: Timestep embedding is inherently prone to vanish
By definition, the time-dependent neural network should yield different behaviors depending on the timestep. The major claim of this study that the architectural choice of the time-dependent neural network affects iime-dependency. In other words, certain incorrect architectural designs disable the time-awareness of the time-dependent neural network.
We first show that the style architecture loses its time-awarenese depending the choice of normal- ization layer. Here, we revisit the ConcatConv operation. We rewrite the convolutional kernel W with the last one W and the rest one WI.c, i.e., W = [WI.o; Wc+1]. Now, the kth channel of the ConcatConv output
z [Wi.cs W [XtJJ -Wh. X + W W+1 +
(2)
(3)
-Wh -W X
(4)
where 4 W
C+1 J lXxx11 Because the last convolution is applied to the all-ones tensor, we obtain vf, = Σ Wk C+1 for all spatial locations (i,j), which is actually single parameter that corresponds to the spatial sum of the kernel (Figure D This interpretation applics to kth channels. From this observation, the ConcatConv operation equivalent to 1) applying vanilla convolution with W = W1.c and 2) subsequently applying clement-wist addition with timestop embodding v, tv, which can be written z W X+ v., Therefore, the pipeline [GN Act. ConcatConv GN Act- ConcatConv GNJ is equivalent to pipeline of [G.N-Act -onv- Eame-Emb--a GN-Act MON-Act-Come Cone----bbbb GN], where the Emb operation indicates adding timestep embedding. Crucially, for each channel, even though the convolutional kernel is parameterized by nine parameters for a 3 3 convolution, its output for the all-ones tensor is simply one parameter, which is equivalent to a scalar offset for a feature map.
The core problem we claim that adding a channel-wise scalar offsct inherently prone to vanish when subsequently applying normalization layer. Before starting the main topic of GN, here consider batch normalization (BN) (Toffe & Szegedy 2015). BN applies whannel-wise mean-std normalization, where the
4

---

Published in Transactions on Machine Learning Research (05/2025)
W
X
Normalize 14,
0
1
0
-1
C/G
Incorporated timestep embedding does not affect the output.
Vt
Normalize 14 o²
1
2
1
2
1
-1
1
W X + Vt
Figure 2: Illustration of vanishing timestep embedding. An additive scalar offset is simply canceled out by the subsequent mean-std normalization.
W * X
Normalize 14, o²
0
1
-1
-1
1
O
1
1
C/G
Incorporated timestep embedding
VE
affects the output.
2
Normalize
1
2
14, o²
-1.4
-1.4
0
2
3
0
1.4
1.4
W
X + Vt
Figure 3: To avoid the use of scalar offset, should ensure that each normalization unit has several elements of timestep embedding more than one in each channel, which would not be canceled out by the subsequent normalization
channel-wise offset is simply canceled out (Figure) Thus, the keh channel of the BN output is
W + X +
E W
X + v
(5)
Var Wh X
(W4 * X W X
(6)
Var W * X
In this example, the timestep embedding v cannot affect the output, which disables dependency on the timestep 1.
5

---

Published in Transactions on Machine Learning Research (05/2025)
In fact, though common neural networks prefer BN, time-dependent neural networks have opted for GN. This practice is in agreement with our analysis: In contrast to BN, which applies channel-wise normalization, GN partitions C' channels into G groups, and each group contains C/G channels that are generally larger than onc. In this case, the C/G elements affect the mean and standard deviation of the normalization, which would not be canceled out (Figure 3). In summary, to avoid vanishing timestep embedding, we should ensure that each normalization unit has the number of channels larger than one. In this regard, GN and layer normalization (LN) (Ba et al., 2016) are safe to be deployed in time-dependent neural networks because they have C/G and C channels per normalization unit, respectively. However, BN and instance normalization (TN) (Ulaanov et al. 2016) have single channel in a normalization unit, which disables timestep embedding.
Note that although the existing practice of employing GN with G = 32 partially solves this problem, we further claim that their timestep embeddings are inherently prone to vanish owing to this architectural choice and still lack sufficient time-awareness. For example, we later claim that the additive timestep embedding is still prone to vanishing, even with GN, depending on the relative variance, whose fundamental cause lies in the architectural configuration. The objective of our study to revisit the architectural choice of time- dependent neural networks, providing practical guidelines beyond choosing G = 32 to further enrich their timo-awareness.
Now we visit, the diffusion models. To discriminate between noise levels that have been applied in a diffu- sion process, diffusion models require time- awwreness for each state. The original study (Ho et al. 2020) implemented a denoising diffuusin probabilistic model (DDPM) using a pipeline of [GN Act Conv Emb GN-Act-Corin Conv]. which is significantly similar to the pipeline of the NODE-style architecture. Since its in- troduction in DDPM, this pipeline has been widely employed in other generative models, such NCSN
(Song et al. 2021ch, DDIM (Song et al. 2021a). VDM (Kingına et al. 2021), and Improved DDPM (Nichol
2021 (See the Appendix). Similar to the NODE-style pipeline, because normalization is applied after additive timestep embedding, we claim that the DDPM-style pipeline is similarly prone to vanishing timestep embodding. Notably, the timestep embedding in diffusion models is obtainod by applying sinusoidal timestep embedding to a multilayer perceptron (MLP) (Ho el al. 2020 Vaswani el al. 2017 Mildenhall al., 2022), which is different from the linear-time NODE-style timestep embedding V, nusoidal timestep embedding in diffusion models may seem advanced, they still have only single clement in each channel, which is vulnerable to being canceled out by the subsequent normalization layer. Thus, we claim that the vanishing timestep embedding can also be observed in diffusson models, which should be prevented in advance. This analysis explains why GN has been preferred in diffusion models.
& Dhhriwal
Although the si-
Note that diffusion models also deploy other operations such as attention blocks. Nevertheless, we find that timestep embedding is only injected in the above pipeline with convolutional layers. Based on current standard practice, our study focuses on the pipeline with convolutional layers.
Summary of the Problem Statement
The timestep embedding in NODE and diffusson models is prone to being canceled out due architec- tural configuration, especially due to the subsequent normalization layer. We refer to this phenomenon as vanishing timestep embedding. Our claim is that the current architectural configurations for NODE and diffuuson models readily cause the vanishing timestep embedding, which should be addressed via architectural correction.
3 Solutions
3.1 Avoid channel-wise scalar offset: Positional timestep embedding
One perspective on the cause of the vanishing timestep embedding is the addition of the channel-wise scalar offset, which is simply canceled out by the subsequent mean-std normalization. To prevent the vanishing timestep embedding, we introduce a spatial degree of freedom in each channel. Specifically, we propose injecting positional timestep embedding:
- W + X V. Pr.
(7)
6

---

Published in Transactions on Machine Learning Research (05/2025)
W
X
Normalize J, o²
0
1
0
1
-1
-1
1
C/G =
V
Incorporated timestep embedding
affects the output.
1
2
3
4
P
Normalize H, o²
4
4
6
-1.4
1.4
W X + Ve
+
Figure 4: Injecting positional timestep embedding enables spatial degree of freedom, which not canceled out by the subsequent normalization
In contrast to the existing timestep embedding v, that has a single element in each channel, the positional timestep embodding p. has H x W clements, which are spatially distinct but shared across all channels (Figure Adding this positional timestep embedding enables its different values a normalization unit, which prevents vanishing timestep embedding, even with channel-wise normalization.
To be compatible with the sinusoidal timestep embedding used in diffusion models, also propose injecting positional timestep embedding using another parallel branch of MLP from the sinusoidal timestep embedding (Figure 5). Specifically, we generate different frequencios for each position and concatenate their sine and cosine to build the sinusoidal, which is subsequently fed to MLP to produce both the existing timestep embedding and the positional timestep embedding. One may choose a lincar-time design positional timestep embedding P. /p or an advanced one using MTP with sinusoidal timestep embedding, where we found that both designs are effective preventing the vanishing timestep embedding.
3.2 Control relative variance: Zero bias initialization convolutions
Another perspective on the cause of the vanishing timestep embedding v, is its small variance compared with the variance in the operand, i..,, Wx X. Because GN applied to the sum of the two, the contribution of the timestep embedding is dependent
their relative scale. If the operand far more dominant than the timestep embodding, then the contribution of the timestep embedding vanishos. Indeod, Kim et al. (2023) proved that when applying LN to the sum of input X and operand P LN(X + P), the gradient (LN(X+P) decays with factor of /Var[X], which means that larger scale of X reduces the contribution DP
of P. This theory applies to our analysis: For GN(W X + vt), to encourage alive timestep embedding, we should relatively strengthen the variance of the timestep embedding v, and reduce the variance of the operand W X.
Indeed, we find that the timestep embedding is inherently prone to vanishing from initialization. Imagine the convolutional kernel w [W1..c; Wc+1] is subjected to initialization method such He initialization He et al. 2015), with zero mean and a variance of o². Note that convolution computes a summation across elements, including the channel unit and kernel with a spatial size ny x May- The input feature map X is the output from the pipeline [GN Act], whose variance is 1/2 at initialization due to the property of the ReLU (He al. 2015). Thus, the variance of the operand W is
C
Var I.W.. X
(8)
2

---

Published in Transactions on Machine Learning Research (05/2025)
Shared layers
Sinusoidal
[FC-Ac-FC-AcC]
Embedding
Private for each layer
FC
FC
Reshape
Reshape
H
W
+
X
Vt
Figure 5: Diffusion models compute sine and cosine from different frequencics and positions, which are fod MLP to produce timestep embedding V,. We propose adding another branch to obtain positional timestep embedding p, from the sinusoidal.
Table 2: Simulation results measure empirical ratio for the relative variance
C
Var WI.C * X
Var|Wo4 C+1 * [J
Relative Variance
64
0.154
4.509 10 10
34.041
128
0.161
2.863 x 10-3
56.089
256
0.166
1.399 x 10
118,869
512
0.164
6.489 x 10-4
252.391
1024
0.168
3.090 x 10
543.949
2048
0.167
1.580 10-4
1056.997
4096
0.166
8.223 10-5
2022.146
By contrast, the timestep embedding is computed from single channel, whose summation includes only the kernel unit of ni Rww Thus, we have
Var C+1 [J) o
(9)
In summary, the relative variance between the operand and the timestep embedding at initialization is
C
Var X//Var
(10)
212
In practice, this ratio is quite large bocause the number channels typically a larger number, such 128. In other words, the timestep embedding inherently exhibits smaller variance compared with that of the operand. Note that the term C here comes from the input to ConcatConv, not from the input to GN. Specifically, the operand branch computes convolution with an input feature that has C' channels; when the number of channels for imput feature C is larger, the variance of the operand becomes higher, which affects the relative variance. In other words, the term C here comes with relative ratio wetween the operand and the timestep embedding, which is quantified by the ratio of the number of channels at initialization.
We also simulated the ConcatConv module initialization and measured its empirical variance (Table 2 For this simulation, set 1, Bh The 3, H W 32, C {64, 128. ,4096), and measured
8

---

Published in Transactions on Machine Learning Research (05/2025)
tinestep
timestep wio timeslop
Test socursey CIFAR-10
Test accuracy (%) CIFAR
82
68
67
80
0.001
0.01
0.1
10
1000 10k
0.001
0.01
0.1
10 100
1000
10k
Initialization for weight
Initialization for weight
Figure 6: Controlled experiments with different initializations W =W1.C that corresponds to the operand branch with C channels, while keeping Wc+1, or equivalently that corresponds to the last channel or equivalenty timestep branch. The results were obtained from the baseline of the NODE setup using group normalization Section 4.1, but with different initializations on weights from a standard deviation of 10 10 The r-axis represents the standard deviation of weight initialization.
variance across 2048 samples. We observed that the measured ratio of relative variance suitably matches with the theoretical expectation of C/2t?, which verifies the above analysis and our claim that timestep embedding is inherently prone to vanish.
Note that we are not saying to increase the variance of the timestep embedding v, as much as possible; our claim is that we should at least avoid a dead timestep embedding with few contribution to the output. This objective can be achieved in a varicty of ways. For example, increasing the initialization scale in v and decreasing the initialization scale in W would be valid. We empirically observed that NODE yielded decreased accuracy in image classification if the initialization scale in W was larger than 102, which dominates over the timestep embedding (Figure 6 Specifically, the performance of those cases was even similar to that of NODE without timestep embedding, which also exhibited degraded performance due to the effect of weight initialixation. Applying a smaller ializatio scale in W leads to alive timestep embedding, which is beneficial to ensure sufficient nime-dependency of NODE. Albeit the validity of this approach, however, discovering the proper initialization scale requires hyperparameter tuming depending on the target tasks.
Alternatively, we propose achieving this objective by modifying the initialization biases. Similar to diffusson models, imagine two branches for the operand via convolution and timestep embedding via MLP, resulting in two distinct bias terms. Note that when there are multiple elements, different bias values within normalization unit cannot be canceled out by subsequent normalization. In this regard, propose 1) applying zero bias initialization for convolutions that correspond to the operand, and 2) applying nonzero bias initialization for timestep embedding, allowing default initialization in Torch He initialization (He et al. 2015). Albeit simple, this approach satisfies the objective and works suitably in practice. For the diffusion- style timestep embedding, this approach indicates applying nonzero bias initialization to the MLP. Indeed, the common practice time-dependent neural networks to use a consistent initialization for all layers using either zero bias or nonzero bias initialization, whereas we propose adopting different. initializations for timestep embedding and operand layers.
3.3 Increase degree of freedom in normalization unit: Decrease the number of groups
As discussod above, the channel- normalization has a single timestep embodding in each normalization unit, which simply canceled out. This scenario corresponds choosing G C' in the GN and is equivalent to IN. When choosing G C/2, each normalization unit has two different timestep embeddings, which would
9

---

Published in Transactions on Machine Learning Research (05/2025)
Table 3: Test accuracy (%) on CIEAR datasets. The standard deviation (Std.) across five runs presented on the right. and indicate disobeying and following our guidelines, respectively.
CIFAR-10
CIFAR-100
Architecture
Accuracy
Std.
NFE
Accuracy
Std.
NFE
Baseline (GN)
91.76
(0.24)
26.0
69.47
(0.54)
26.0
Replace GN with BN
91.11
(0.31)
34.6
68.01
(0.40)
71.3
+ w/o timestep embedding
91.06
(0.23)
28.7
68.09
(0.22)
69.6
Inject P
91.80
(0.44)
26.0
69.86
(0.76)
26.0
not be canceled out but are still prone to vanish owing to the lack numbers. Choosing a sufficiently smaller G increases the number of different timestep emboddings C/G in each normalization unit, which safeguards against the vanishing timestep embedding. In fact, even though choosing G 32 the de facto standard, few implementations of diffusion models such Score SDE (Song et al., 2021ch, ScoreFlow (Song et al., 2021b), and DPM-Solver (Lu et al. 2022) specify G = min(C/4432 in their source code, which restricts the lower bound C/G 2 4. Despite its validity, there has been no explanation for this source code in their paper, nor validation of the lower bound of four. Later, we empirically examine the effect of the number of groups. Indeed, choosing G = yielded significantly degraded performance, whereas choosing the counterpart in the extreme case of
yielded the most salisfactory results.
4 Experiments
We examine the effect of the three proposed methods on solving the vanishing timestep embedding, targeting NODE and diffusion models.
4.1 Experiments on NODE
The target model was NODE for image classification. The early-stage consists of one 3 3 convo- lution and two residual blocks with a stride of 2, which produces an initial value for the ODE. The NODE produces output to which head is applied. The head is pipeline of [GN Act GAP Dropout FC Softmax], where GAP indicates the global average pooling layer. We used ReLU activation function, the mumber of channels C' 256, dropout (Srivastava et al. 2014) with probability of 0.5, and the number of groups G 32 unless specified otherwisc. For the ODE, we used relative and absolute tolerances of 0.001 and solver of Ktunge-Kuta (Runge 1895: Kutta 1901 of order 5 of Normand-Prince-Shaane
Model
(dopri5) (Dormand & Prince 1980 with the adjoint method, implemented in PyTorch (Chen et al. 2018
Hyperparamters The target datasets were CIFAR-(10, 100}, which consist of 60K images of {10, 100} classes [Krizhossky et al. 2009). For data augmentation, we used 32 x 32 random cropping with 4-pixel padding, color jitter with hue degree of 0.05 and saturation degree of 0.05, a random horizontal flip with a probability of 0.5, and mean-std normalization using dataset statistics. For training, the number of epochs of 250, stochastic gradient descent with a momentum of 0.9, learning rate of 0.1, learning rate decay that reduces the learning rate when metric has stopped improving decay of 0.0001, and mini-batch size of 128 were used. In fact, early studies on NODE used a lower benchmark and reported a test accuracy of 60% for CIFAR-10 (Dupont et al. 2019; Xia et al. 2021), but we targeted a more docent benchmark with a test accuracy above 90% for CIFAR-10, which similar or even better to the recent results of Kim et al. (2024). An average five runs with different random seeds was reported. The training was conducted on single RTX 3090 GPU machine.
platean with a patience of 15, weight
Do not use BN. We first examine the effect architectural choice. The baseline performance here was obtained with a NODE-style pipeline with ConcatConv. We observed that NODE yielded decrcased accuracy if GN was replaced with BN (Table This observation demonstrates the suboptimal performance of time-dependent neural network due to the channel-wise normalization of BN. We also report results
10

---

Published in Transactions on Machine Learning Research (05/2025)
Table 4: Test accuracy (%) on CIFAR datasets for different bias initializations
CIFAR-10
CIFAR-100
Zero Bias Initialization
Accuracy
Std.
NFE
Accuracy
Std.
NFE
Baseline (All)
92.03
(0.31)
26.0
69.36
(0.57)
26.0
Timestep branch
92.00
(0.33)
39.2
68.91
(0.79)
26.0
Convolution branch"
92.18
(0.31)
26.0
69.57
(0.68)
26.0
with BN but without timestep embedding. Indeed, one may think that NODE without timestep embedding would not work suitably, such as with test accuracy of 10%; nevertheless, we observed that it achieved a suboptimal but quite high test accuracy around 90%. This observation is intriguing because NODE without a timestep embedding is something like a repeating single residual block. We conjecture that the repetition of a single NODE block as well as using the last fully connected layer would be sufficient for basic feature extraction and classification when using rich regularizations such as dropout and weight decay. Furthermore, NODE without timestep embedding can be understood as fexgh- RNN with a residual connection, which is already capable of universal approximation, thereby ensuring reasonable performance (Schäfer & Zimmermann 2006). Note that even when using BN, NODE could exploit other factors such as padding. If padding is applied in ConcatConv, the edges of the timestep embedding exhibit different values, which would not be canceled out by subsequent normalization. However, the edge effect of padding was still ive enough to fully solve this problem owing to its suboptimal performance practice, which was similar to the performance without limestep embedding. This result supports our claim that the vanishing of timestep embedding due to BN is dominant in practice, and other indirect factors, such as padding, affect minor. Furthermore, it significantly increased the number of function evaluations (NFE) due to the ineffectiveness of the edge effect in incorporating timestep, which required more computational costs. Therefore, we should avoid BN and adopt other proper methods to prevent the vanishing timestep embedding.
Positional timestep embedding improves performance. We injected the positional timestep embed- ding into the baseline above. For both CIFAR-10 and CIFAR-100, the NODE with positional timestep embedding demonstrated improved accuracy compared with the baselines (Table 3 These results were obtained with the same level of NFEs, which indicates that injecting the positional timestep embedding did not harm solving the ODE.
Apply zero bias initialization to the convolution branch. We examine the effect of bias initializa- tion. In contrast to previous experiments, the bascline performance here was obtained by a NODF-style pipcline with additive timestep embodding using MLP to deploy two branches with distinct bias terms, where we applied zero bias initialization to both branches of timestep and operand, i.e., convolution. We further obtained results from other setups of applying zero bias initialization to one branch and nonzero bias initialization to another branch. We observed that applying zero bias initialization to the timestep branch slightly degraded performance, whereas applying zero bias initialization to the operand branch improved the test accuracy (Table 4). These results are in agreement with our analysis: To encourage alive timestep embedding, we should apply zero bias initialization to the operand branch.
Decrease the number of groups. We examine the choice of the number of groups in GN. Targeting
two cases of the number of channels 256 and C== 64, we set the number of groups G E {1, 4, ,C} and evaluated the test accuracy. Figure 7 summarizes the results. As the number of groups increased, the test accuracy dropped. This observation is in agreement with our claim. For example, if G = C = 256, the timestep embedding acts as an additive scalar offset, which is canceled out. To avoid this phenomenon, we should ensure a sufficient number of parameters C/G in cach group by selecting a smaller number of groups. This analysis applies to the extreme case of G = 1, which enables the richest parameters in a normalization unit. Indeed, we observed that. G 1 yielded the best. performance in most cases.
11

---

Published in Transactions on Machine Learning Research (05/2025)
C-256
C=256
C-04
accuracy
accuracy
60
87
86
2
32
128
256
16
32
258
The number of groups
The number of groups
Figure 7: Results for different choices of the number of groups on CIFAR-10 (left) and CIFAR-100 (right). The average and standard deviation from the result of five random seeds are depicted.
Table 5: Results of the DDPM experiment with step-by-step validation of the proposed method
Setting
Inject Pt.
Zero bias
Set G
1
FID (4)
IS (t)
Baseline
x
X
3.238
9.507
+ Inject Pr
/
X
3.199
9.539
ZZero bias
/
X
3.122
9.549
+Set G =
/
/
3.074
9.603
4.2 Experiments Diffusion Models
The target model DDPM, which is the standard diffusion model for image generation. Because main larget is to confirm the vanishing limestep embedding in existing diffusion models, rather than targeting custom model, we used the same architecture of existing DDPM. The model uses a U-Net architecture (Ronneberger et al. 2015) with sinusoidal timestep embedding and consists of residual (He et al.) 2016) and sell-attention blocks Vassani et al. 2017). We used Swish activation function (Ramachandran et al.) 2018 Elfwing et al. 2018), model width of 128, dropout probability 0.1, and the number of groups G = 32 unless specified otherwise. For the diffusion process, Pmax = 20 and 3min 0.1, and noise scales with mmin 0.01 and max = 20 with total number of 1000 were used.
The target dataset was CTFAR-10 with 32 x 32 resolution, where a random horizontal flip with a probability of 0.5 was applied for data augmentation. For training, the number of iterations of 1300K, Adamn optimizer (ingmma & Ba, 2015) without weight decay, learning rate of 0.0002, gradient clipping (Pascanu et al. 2013) with maximum norm of 1, warmup steps of 5K, and mini-batch size of 128 were used. The training was conducted on a 4x RTX 3090 GPU machine, which required approximately seven days for training and sampling.
Table 5 summarizes the results. The three proposed methods improved both the Fréchet inception distance (FID) (Huusel et al.] 2017) and the inception score (IS) (Salimans et al. 2016). The improvements were observed through a step-by-step application, which demonstrates the validity of our methods. Note that a common approach to architectural modification, such as stacking more layers, improves FID and inception score but with higher inercased computational complexity, such as the number of parameters and FLOPS. However, our proposed methods bring performance gains without largely affecting computational complexity, which we believe is reasonable improvement. Indeed, the main architecture of U-Net occupies most of the
12

---

Published in Transactions on Machine Learning Research (05/2025)
number of parameters and FLOPS, and the proposed modifications of timestep embeddings affect little these computational complexities while solving the critical problem that is inherent to the architecture.
5 Discussion
Related works To our best knowledge, the vanishing limestep embedding is first discovered by this study, However, existing literature might have unknowingly taken advantage of partially solving this problem. For example, the study on ANODE (Dupont et al.] 2019) proposed augmenting an ODE space to improve the representational limitation of NODE. From our analysis, a larger dimension of the ODE increases the number of timestep embedding elements C/G in each normalization unit, which prevents the vanishing timestep embedding. For G = 32, when augmenting the dimension of the ODE from 64 to 128, the number of timestep embedding elements changes from two to four, where the latter is more stable against vanishing timestep embedding. However, our theory says that the same effect can be achieved by simply reducing the number of groups from 32 to 16. Furthermore, the larger dimension of ANODE requires a higher computational cost and an increased number of parameters, which is an inefficient approach to solving the vanishing timestep embedding compared with our methods, albeit the validity of augmenting the ODE space to address the representational limitation, Note that the notion of same effect here indicates the equivalence in C/G unit regarding the factor of vanishing timestep embedding, not the equivalence in performance. Figure 7 shows that the models with C 256 achieve improved performance overall compared with that of C= 64. In other words, increasing the number of channels C' would enlarge the width of the neural network, which itself is advantagcous to improve the performance.
In fact, the architecture of vanilla NODE has been tackled by several studies [hhan el al. 2019 Avelin & Nyström 2019 Massaroli et al. 2020: Queiruga et al. 2020 Choromanski et al. 2020). They have mentioned that the continuous-time version of residual networks should be incarnated by time-dependency in paramcters 0. but the vanilla NODE rather implements dimeddepnnyy in f through the design of the neural network. Despite the validity of ensuring time-dependency in parameters 0, the current practice (Dupont et al. 2019; Yan et al. 2020 Norcliffe et al. 2020 Kidger et al. 2021 Xia et al., 2021) still prefers the standard pipeline with ConcatConv owing to its simplicity, and this pipeline has even influenced the implementation of diffusion models. Our study advocates for both theory and practice: The proposed methods are compatible with the cxisting pipeline of NODE and diffusion models whilc cnriching their insufficient time-dependency.
Also, in the main text, we focused on standard diffusion models, which adopt the DDPM-style pipeline. However, some of the recent diffusion models such as EDM2 (Karras et al.) 2024) adopt a modified pipeline, which applics multiplication of the timestep embedding rather than the existing pipcline with the addition of timestep embedding. Because this modified pipeline does not apply subsequent GN, our analysis says that it is indeed desirable to prevent the vanishing timestep embedding. Nevertheless, we find that their study improved performance without providing a convineing explanation for the underlying cause of the performance improvement. Here, our study enables us to determine whether a timestep embedding would be alive or vanish, which is expected to be generic criterion for valid architectural design.
Notes on the known issue of increased NFE The dopriš solver, which is commonly used in NODE, corresponds to solvers with adaptive step size, which choose a large step in smooth dynamics and a small step in varying dynamics. It controls the step size by comparing the predefined tolerance with the current error. For solving NODE, there has been a known issue with increased NFF, which incurs longer computational time. In the official repository of NODE, several users have reported excessively large NFEs in certain usages. The authors of NODE conjectured that the use of the ReLU activation function might yield a stiff ODE due to its discontinuity. This behavior results in a zero step size and an infinite NFE for solving the ODE. They mentioned that when facing this issue, NODE should be modified to use another continuous activation function such as Softplus, which approximates ReLU with a continuous function.
We find that this approach is partially valid but requires correction. Indeed, we empirically observed de- creased NFEs when using continuous activation functions such as Softplus and Sigmoid, at the expense of decreased test accuracy (See the Appendix). However, we find that increased NFE can be caused by both stiff and smooth surfaces; when the surface is too smooth such as plateau, moving step cannot reach
13

---

Published in Transactions on Machine Learning Research (05/2025)
desired objective. We observed this phenomenon of increased NFE for NODE with ConcatConv and BN (Table where the timestep embedding indirectly affects the output through the edge effect of padding. To stabilize the ODE, we should ensure that the surface is neither too stiff nor too smooth by regulating an adequate amount of time-dependency.
Therefore, when facing increased NFE, we should confirm whether the problem caused by large or small stiffness and approach it either by decreasing or increasing stiffness into its stable range. The problem of vanishing timestep embedding corresponds to the latter, and thus our proposed solutions, such injecting positional timestep embedding, can be solution to the increased NFE problem.
6
Conclusion
This study reported the vulnerability of vanishing timestep embedding in NODE and diffusion models. We claimed that in their architectural pipeline, timestep embedding is prone to vanish when subsequently applying normalization. To address this problem, we proposed 1) adding positional timestep embedding, 2) applying zero bias initialization in the convolution branch, and 3) reducing the number of groups. The proposed methods consistently improved the performance NODE and diffusion models, which indicates that the two families of models inherently pose this problem, and encouraging alive timestep embedding using these explicit methods is significantly beneficial for their current implementations.
Acknowledgments
This work was supported by Samsung Electronics Co., Ltd (10201210-08019-01), Also, this work was supported by JST CREST Grant Number JPMJCR1913.
References
Benny Avelin and Kaj Nyström. Neural ODEs as the Deep Limit of ResNets with constant weights. CoRR, abs/1906.12183, 2019.
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization. CoRR, abs/ 1607.06450. 2016.
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural Ordinary Differential Equations. In NeurIPS, pp. 6572-6583, 2018.
Krzysztof Marcin Choromanski, Jared Quiney Davis, Valerii Likhosherstov, Xingyou Song, Jean-Jacques E. Slotine, Jacob Varley, Honglak Lee, Adrian Weller, and Vikas Sindhwani. Ode to ODE. In NeurIPS, 2020.
John Dormand and Peter Prince. A family of embedded Runge-Kutta formulae. Journal of computational and applied mathematics, 6(1):19 26, 1980.
Emilien Dupont, Arnaud Doucet, and Yee Whye Teh. Augmented Neural ODEs. In NearIPS, pp. 3134-3144 2019.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
linear units for neural network function
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving Deep into Rectifiers: Surpassing Human- Level Performance on ImageNet Classification. In ICCV, pp. 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. In CVPR, pp. 770-778, 2016.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs Trained by Two Time-Scale Update Rule Converge to Local Nash Equilibrium. In NIPS, pp. 6626- 6637, 2017.
14

---

Published in Transactions on Machine Learning Research (05/2025)
Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising Diffusion Probabilistic Models. In NeurIPS, 2020.
Sergey loffe and Christian Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. In JCML, volume 37, pp. 448-456, 2015.
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and Improving the Training Dynamics of Diffusion Models. In CVPR, pp. 24174-2411 2024.
Patrick Kidger, Ricky T. Q. Chen, and Terry J. Lyons 'Hey, that's not an ODE': Faster ODE Adjoints via Seminorms. In JCML, volume 139, pp. 5443 5452, 2021.
Bum Jun Kim, Hyeyeon Choi, Hyconah Jang, Dong Gu Lee, Wonscok Joong, and Sang Woo Kim. Improved robustness of vision transformers via prelayernorm in patch embedding. Pattern Recognit., 141:109659, 2023.
Semin Kim, Jaehoon Yoo, Jinwoo Kim, Yeonwoo Cha, Saehoon Kim, and Seunghoon Hong. Simulation- Training of Neural ODEs on Paired Data. CoRR, abs/2410..22 2024.
Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In ICLR, 2015.
Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational Diffusion Models. In NeurIPS, 2021.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
Wilhelm Kutta. Beitrag zur näherungsweisen Integration totaler Differentialgleichnngen 1901.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps. In NeurIPS, 2022.
Stefano Massaroli, Michael Poli, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama. Dissecting Neural ODEs. In NewrIPS, 2020.
Ben Mildenhall, Pratul P. Srinivasan, Matthew Tancik, Jonathan T. Barron, Ravi Ramamoorthi, and Ren Ng. NeRF:: representing scenes neural radiance fields for view synthesis. Commun, ACM, 65(1):99 106, 2022.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved Denoising Diffusion Probabilistic Models. In ICML, volume 139, pp. 8162-8171, 2021.
Alexander Norcliffe, Cristian Bodnar, Ben Day, Nikola Simidjievski, and Pietro Lió. On Second Order Behaviour in Augmented Neural ODEs. In NearIPS, 2020.
Razvan Pascanu, Tomás Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In ICML (3), volume 28, pp. 1310 1318, 2013.
Alojandro F. Queiruga, N. Benjamin Erichson, Dane Taylor, and Michael W. Mahoney. Contimuous-in-Deppl Neural Networks. CoRR, abs/2000 /2//220080038) 2020.
Prajit Ramachandran, Barret Zoph, and Quoe V. Le. Searching for Activation Functions. In ICLR (Work- shop), 2018.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional Networks for Biomedical Image Segmentation. In MICCAI (3), volume 9351, pp. 234-241, 2015.
Carl Runge Über die numerische Auflösung von Differentialgleechnngn. Mathematische Annalen, 46(2): 167 178, 1895.
Tim Salimans, Ian J. Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved Techniques for Training GANs. In NIPS, pp. 2226 2234, 2016.
15

---

Published in Transactions on Machine Learning Research (05/2025)
Anton Maximilian Schäfer and Hans-Georg Zimmermann. Recurrent Neural Networks Are Universal Ap- proximators. In ICANN (1), volume 4131, pp. 632-640, 2006.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising Diffusion Implicit Models. In ICLR, 2021a.
Yang Song, Conor Durkan, Tain Murray, and Stefano Ermon. Maximum Likelihood Training of Diffusion Models. In Newr/PS, pp. 1415 1428, 2021b.
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. In ICLR, 2021c
Nitish Srivastava, Geoffrey E. Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting, J. Mach. Leorn. Res., 15(1):1929-1955 2014.
Dmitry Ulyanov, Andrea Vedaldi, and Victor S. Lempitsky. Instance Normalization: The Missing Ingredient for Fast Stylization. CoRR, has/1607.08022 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In NIPS, pp. 5998 6008, 2017.
Yuxin Wu and Kaiming He. Group Normalization. In ECCV (13), volune 11217, pp. 3 19, 2018.
Hedi Xia, Vai Suliafu, Hangjie Ji, Tan M. Nguyen, Andrea L. Bertozzi, Stanley J. Osher, and Bao Wang. Heavy Ball Neural Ordinary Differential Equations. In NeurIPS, pp. 18646-18659, 2021.
Hanshu Yan, Jiawei Du, Vincent Y. F. Tan, and Jiashi Feng. On Robustness of Neural Ordinary Differential Equations. In ICLR, 2020.
Tianjun Zhang, Zhewei Yao, Amir Gholami, Joseph E. Gonzalez, Kurt Keutzer, Michael W. Mahoney, and George Biros. ANODEV2: A coupled neural ODE framework. In NeurIPS, pp. 5152 5162, 2019.
A Results on Other Activation Functions
Table summarizes experimental results on other activation functions for NODE. When replacing ReLU with a continuous activation function, obtained decreased NFE, which leads faster computation. For the Sigmoid function, obtained an NFE of 14.0, which is approximately half of the NFE of 26.0 for ReLU, Similar results were confirmed for Softplus, ELU, and SiLU activation functions. Albeit the advantage of decreased NFE, the continuous activation functions yielded decreased Lest accuracy.
Table 6: Results other activation functions
CIFAR-10
CIFAR- 100
Activation Function
Accuracy
NFE
Accuracy
NFE
ReLU (Baseline)
91.76
26.0
69.47
26.0
SiLU
91.76
21.5
68.72
26.0
ELU
91.38
23.5
67.93
26.0
Softplus
89.89
18.8
65.88
16.0
Sigmoid
89.75
14.0
65.43
14.0
B Current Implementations of Diffusion Models
The pipeline of CON-Act-Conv-Emi Conv] has been widely used in current implementations of diffusion models. Here, provide several examples.
16

---

Published in Transactions on Machine Learning Research (05/2025)
Reference: https:// github. con/hoj wwwtoannooddttteeeeeaaeeaaooaeee diffuiii def resnet_blockxx, temb, name, out_ch-None, cony_hhortcut-Falsee, dropout):
=f.varisble__cae scope (name):
nonlinearity (normalize (h, temb=temb, name='norni")
h nn. convdd name='convi", nun nun_units=out_ch)
10
# add in timeatep embedding
11
h += na dense www.eeeennninnnnenne (nonlinearity (temb), tame='tenb_proo aun_niis==out_ch) [:, None, None
:]
h nonlinearity (normalizoo (h, temb=tenb, name= lame='now norm2
Listing 1: Implementation of DDPM
Reference: bttps://githbb sotpaa/eooooocoooseeeeeeeee class RosnetBlockDDPM (nn (od.ullee)))
#
def forward
temb=None):
#
h self act (self. GroupNorm_O(x))
h self. Conv_O(h)
Add bias to each feature map conditioned the time enbedding temb
None:
self Dense_0 (self. (temb)) [:, :, None, None] self aelf GroupNorm_.
Listing 2: Implementation of Score SDE
Reference: https://githut con/ con//rrongrrup/deee ddim class ResnetBlock (nn. Module):
def forward x, temb);
self normi (h)
nonlinearity (h)
self convi (h)
10
self. temb_proj (nonlinearity (temb)) [:, :, None, None]
11
19
helf.norm2 (h)
13
h nonlinearity (h)
14
Listing 3: Implementation of DDIM
Reference: wtwp://hhh....ooooooooooorereraa
class ReenetBlock (nn. Module):
def (self_ x, cond, deterministic: bool, enc-llone):
h h nonlinearity (nnomaaiityyormmalzeeee (normalizel(h)
h an. Convo featuressout_ch, kernel_size=:3, 3), strides=(1, 1), name= conv1 (h)
# add in conditioning
1f cond is not None:
cond. shape [0] and len (cond, shape)
Dense(
features=out_ch ch, bias=False. kernel, init=nn htt..iitt...iiiillirr...... zeros,
cond_proj !) (cond) [:, None, None,
17
17

---

Published in Transactions on Machine Learning Research (05/2025)
nonlinearity (normalize((b))
Listing 4: Implementation of VDM
18


Reviews:
[{"review_id": "23tnFYn2oi", "reviewer": "Reviewer_TPBA", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The three methods proposed in this paper, including positional timestep embedding, zero bias initialization, and fewer groups, are intuitive and simple, while the performance improvement is not significant. I would be more excited if the author could discuss the solutions with more depth in theory and in a broader scope. For example, how NODE can maintain the time information with heterogeneity of neurons' time constants which can approximate the Laplacian basis.\n\nThe paper overall is well-written and contains sufficient details. I don't have specific criticism on techniques.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes methods for solving the vanishing timestep embedding in ordinary neural differential equations.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See the comments in strengths and weaknesses.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "o5YZmQTw9j", "reviewer": "Reviewer_xR1m", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Simple solutions to address the issue of vanishing tilmestep embedding influence in a neural network modeling a time-varying function. \n2. Helps provide a simple justification to the use of GroupNorm over BatchNorm as well as C/G factor in diffusion models.\n\nWeaknesses:\n1. It’s unclear if in practice, such tilmestep embedding vanishes to such an extent that the model is unable to learn the time-dependent behavior. For instance, in Tab.2, even with batch-norm, it was able to recover very close to the optimal performance. Similarly, In Tab. 4, there’s not a lot of difference between the baseline and the scheme with all three mitigations in terms of the inception score. \n2. Although this work performs experiments on Diffusion models which injects timestep embedding through convolutional layers, many recent and advanced architectures injects the timestep embedding through modulation layers. It’s unclear how the proposed claims/arguments apply to these diffusion models.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work studies the disappearance of tilmestep embeddings in neural networks that model a time-varying function, example neural ODEs and diffusion models. This paper analyzes few neural networks with timestep input and shows that some architecture choices (such as batch-normalization) lead to vanishing time-step embeddings, resulting in loss of tilmestep-aware modeling capabilities.  It also proposes a solution that includes incorporating positional embeddings after the tilmestep embedding to negate the effect of batch-normalization layers. Numerical experiments (for Neural ODEs and Diffusion Models) are performed to strengthen these arguments in NODE and diffusion model training.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I believe answering the following questions would help strengthen this work. \n1. Diffusion models also incorporate tilmestep embeddings through modulation layers as well, which gets injected into each block in the form of a gating mechanism. How does timestep embedding vanishes in such instances?\n2. How does one extend this analysis to transformer networks modeling time-dependent function?\n3. In Tab.2, even the solution with batch-norm achieves close to the optimal performance. Doesn’t it signify that tilmestep embeddings do not vanish to the extend this work shows it does?\n4. In Tab.4, why is it the case that the baseline performance is very close to the performance of the scheme with all mitigations applied (for instance the inception scores are nearly similar)?\n5. For diffusion models, have you tried to explore the effects of timestep disappearance with more depth in the network? Does it worsen this issue?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "13KuJc82bN", "reviewer": "Reviewer_74vv", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths \n- Relevant Topic: The study addresses an important aspect of time-dependent neural networks, which has implications for various applications, including time series modeling and diffusion models.\n\nWeaknesses \n- Unclear Problem Definition: The vanishing time step problem is not clearly defined. The paper lacks a precise formulation of the issue, making it difficult to assess its significance.\n- Weak Justification of Non-Vanishing Time Step: The connection between non-vanishing time step embeddings and final model performance remains vague. It is not convincingly demonstrated whether ensuring non-vanishing time step embeddings is beneficial.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper explores the impact of architectural choices in time-dependent neural networks on their ability to retain time-awareness. The authors argue that the design of time step embeddings significantly affects model performance, particularly in diffusion models.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Clearly Define the Vanishing Time Step Problem: The authors should explicitly state what constitutes the disappearance of time step embedding and why it is a problem. Providing mathematical or empirical evidence would be beneficial.\n\n- Demonstrate the Utility of Non-Vanishing Time Step Embeddings: Instead of assuming its usefulness, the authors should first provide a compelling demonstration that maintaining time step information is desirable in practical applications.\n\n- Establish a Concrete Link Between Time Step and Performance: The paper should clarify how time-awareness (or its absence) affects final performance in image modeling tasks like diffusion models. The relationship should be explicitly demonstrated rather than assumed.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "XMKFYbIKvG", "reviewer": "Reviewer_uB6n", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n- The authors point out a failure mode where timestep embeddings have only a very small contribution to the variance in the normalization layer. If time information is important for the task, such as for complex time-series data, this would be detrimental.\n- This may be a problem that has remained under the radar in some off-the-shelf implementations of timestep embedding in convolution based NODEs and diffusion models.\n- The optimization of positional encoding has played a significant role in improving transformers on NLP tasks, and analogously optimization of time embedding for models trained on time-series tasks will be important.\n\n\n## Weaknesses\n\n- Very general claim in the abstract, not adequately supported by the specialized content of the main text, which is limited to ConcatConv + GN for NODEs and scalar time embeddings in conv diffusion models.\n- The text focuses on \"modern\" networks and current practice, but references only up to 2021.\n- Strongly limited discussion of related work (e.g. nothing on other methods of time embedding, parallels to positional encoding in attention).\n- The validation currently does not establish that the observed improvements in the experiments are due to better time dependency. This is partly because no baseline without time embedding is given, and partly because small differences in accuracy are reported without uncertainty estimates.\n- While prominent in the motivation, the paper does not deal with time-series data, but with static tasks (images) where time plays an indirect role as the depth-coordinate in NODEs and the denoising-flow coordinate in denoising models. \n\nIn sum, the manuscript exposes a relevant issue in how timestep information is embedded in the standard NODE implementation and some diffusion models. However, these findings are not carefully argued and validated, and overly general claims are made. Therefore I believe significant revisions are required. Please find my detailed concerns below.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors argue that the ConcatConv method of embedding time information in convolutional NeuralODEs (NODEs), does not function properly when using GroupNorm (GN) with a large number of groups afterwards. In the extreme case of n_groups = n_channels which is equivalent to using BatchNorm (BN), time information is fully canceled out by the normalization. This can be explained by small (or vanishing) variance of the time embedding contribution relative to the variance of the remaining representation at the normalization step. The issue can also arise in convolutional blocks in diffusion models.\n\nThe authors propose that the time embedding will be more prominent when i) using an embedding with positional variance, ii) zeroing biases to reduce variance not related to time information, iii) reducing the number of groups in GN, and show for validation that these modifications may have small positive effects on accuracy of NODEs for image classification on CIFAR-10 /-100 and on image generation in a diffusion model trained on CIFAR-10.\n\nI am not closely familiar with cutting-edge practice in training NeuralODEs and diffusion models, and will therefore not comment on how widespread the use of ConcatConv + GN is and how relevant the contribution is for practitioners.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "## 1. Clearer presentation and representation of the results\n\n - The abstract is vague and nebulous - it should be rewritten to reflect the concrete findings.  \n - It would also be much better to not use the term \"modern time-dependent networks\", which is ill-defined and could as well apply to State-space models, Transformers or RNNs.\n - The title suggests that the timestep information actually disappears in common/practical cases. However this would only be so if BN or equivalently GN with n_groups = n_channels were used in practice, which apparently is not the case.\n - References to current practice for NODE stop at 2021 and for the diffusion models at 2022. Could one or two more recent references be added at least for NODEs?\n - The main motivation given in the beginning is time-series input. But then the experiments are conducted on CIFAR10/100, so static image classification and generation. What is the use of time-information here, how does it relate to augmenting the flow with additional dimensions (ANODE)? The reasons to focus on static tasks instead of time-series input should at least be explained to the reader.\n - (optional) It would be good to extended and improve the discussion of related works. For example, currently no overview of other methods of constructing the time embedding is given, and there is no discussion of the related and much more broadly studied methods of positional encoding in attention blocks. These are only suggestions.\n - Fig.6: Caption and text where Fig.6 is referenced have no details or crossreference to details on what exactly was the model architecture or other parameters here.\n - What is meant by applying nonzero bias init only for the timestep embedding? In my understanding biases in Conv layers are usually added only after summing the convolution over input channels (see e.g. docstring of conv2d in pytorch), so when the time part is already mixed with the operand part, $ W^k \\ast [X; tJ] + b^k$ where $k$ is the output channel index. So here there would be no separate biases for the time and the operand parts. Or do the authors in their code only add the time part explicitly after the convolution with $\\tilde{W}$ as in eq.4, instead of concatenating $tJ$ before? This could be more clearly stated by adding the biases to the equations.\n\n \n## 2. Improving the line of arguments and validation experiments\n\nThe argumentation from theoretical insight to analysis and solutions is not stringent. Please find my concerns on each part of the argument in the following points.\n - At first, a clear explanation of how BN fully removes the timestep information added by ConcatConv is given. However, then it turns out that BN is never used but GN instead where the problem does typically not arise (and it is not communicated beforehand that the BN discussion is just to introduce to the main topic of GN). Unfortunately, here the theoretical argument breaks off and continues only in anecdotal fashion. The point that ConcatConv results in only a single effective learnable parameter per channel, making the time embedding little expressive, is not expanded upon. The separate issues of small expressivity and of \"vanishing\" are mixed in the discussion (in sect.3 and also end of sect.2). In sect.2, no arguments are given why the timestep information would still be prone to vanish when GN is used. The concept of relative variances appears only later in the Solutions sect.3.\n - In sect.3 the argument is made that the relative variance of the timestep part may be too small compared to the operand part carrying the sample features. However, the size of these variances is neither analyzed theoretically, nor measured empirically, although both could easily be done (at init for the theory). Nor is any argument given that it is the size of the variance, and not the expressivity of the embedding, which makes the timestep information insufficient. For these reasons, this argument remains in a state of motivation, not of concrete explanation.\n - If small relative variance of the timestep part is detrimental, conversely at the other extreme also small relative operand variance should intuitively become detrimental or result in longer training times. Why is this not observed in Figure 6, where the accuracy remains high also for the smallest W initialization scale? Also, at what point on the x-axis in Fig.6 is the init variance of the operand part equal to that of the time embedding part?\n- Fig.6: The drop in performance for large W init scale could well be due to issues with the gradients unrelated to time-embedding. To demonstrate that the effect is due to varying the relative variances of time and input features, a baseline should be shown in the figure where no time embedding was used at all, while the init scale of W is varied as in the experiment.\n- The accuracy numbers reported in Tables 2 and 3 are averages across 5 runs. Since the differences in accuracy especially between baseline and the proposed positive intervention is small, the standard error of the mean should be reported (and it might be necessary to reduce it by increasing the number of runs).\n- In 4.1, par. \"Do not use BN\", it is argued that due to padding also with BN some time information enters the network (although with increased NFE), and that this is the reason why the performance would remain relatively high. To make this point properly it would be necessary to provide a baseline where no time embedding is used at all. Would this cause a much larger reduction in accuracy?\n- In related work discussion: The authors claim that the ANODE benefit from doubling the internal dimension might be achieved also by halving the number of groups in the GN operation, if the benefit can be explained by making the time-embedding more important. This experiment could simply be performed, and would significantly strengthen the current paper, either by highlighting the importance of the time embedding compared to the dimensionality of the representation, or else by putting the results into a clearer perspective. \n\n\nOverall, please do not see the number of my concerns as aimed at preventing eventual publication of the paper in TMLR, by possibly requiring the authors to decrease their claims. I strongly believe results only need to be clearly communicated and valid to warrant publication, regardless of hypothetical impact. In this respect also TMLR has no impact requirement, only that some readers would be interested in the results, which I believe will be the case.", "broader_impact_concerns": "/", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers but one lean towards acceptance of the paper.

The consensus among reviewers is that the issues underlined by this paper are real, but they might not be so severe in practice due to the discrepancy between the specific layers studied in this work and the layers used by practitioners. Furthermore, the fact that the improvement observed in the experiments is due to the fixation of the issue of disappearance ot time embedding could be clearer.

Still, all reviewers found the message of this work interesting, that it will interest several researchers and maybe pave the way to a better understanding of timestep embeddings in recent neural networks.
Decision: Accept (poster)