TD-LN: Time-Dependent Layer Norm

• TD-LN is a normalization method that adapts parameters based on time-indexed statistics, addressing static limitations in sequential and generative models.
• It utilizes two principal approaches: spatio-temporal aggregation (ATN) over time windows and time-parameterized gain/bias to integrate temporal information.
• Empirical evaluations show that TD-LN improves performance and efficiency, reducing losses in language tasks and lowering FID scores in diffusion-based architectures.

Time-Dependent Layer Normalization (TD-LN) refers to a family of normalization techniques in neural networks where the normalization statistics or parameters explicitly depend on the time index, in contrast to conventional Layer Normalization (LN) which is inherently time-invariant. TD-LN addresses key deficiencies of LN in sequential and generative modeling settings by leveraging temporal or time-step information within normalization. Two principal methodologies have emerged: (1) “spatio-temporal” normalization, which aggregates statistics across contiguous time windows (exemplified by Assorted-Time Normalization [ATN]), and (2) time-parameterized rescaling and shifting, which replaces static LN gains and biases with time-conditioned versions, notably in diffusion models. These approaches yield theoretical and empirical improvements in recurrent networks and diffusion-based generative architectures over time-invariant normalization schemes (Pospisil et al., 2022, Liu et al., 13 Jun 2024).

1. Motivation for Time-Dependence in Layer Normalization

Standard Layer Normalization computes feature-wise mean and variance for each time step individually. For RNNs, the per-time-slice mean $\mu_t$ and variance $\sigma^2_t$ are

$$\mu_t = \frac{1}{n}\sum_{i=1}^n a_i^{(t)}, \quad \sigma^2_t = \frac{1}{n}\sum_{i=1}^n (a_i^{(t)} - \mu_t)^2$$

where $a^{(t)}\in\mathbb{R}^n$ is the pre-activation vector at time $t$. LN normalizes via

$$y^{(t)} = \gamma \odot \frac{a^{(t)} - \mu_t}{\sqrt{\sigma^2_t+\epsilon}} + \beta$$

with static trainable $\gamma, \beta \in \mathbb{R}^n$. This “washes out” temporal drift in input norms or means, preventing information about gradual, task-relevant shifts in the sequential input from propagating through the network. Furthermore, the invariance of LN to rescaling at each step blocks the RNN from encoding norm-based temporal signals. Similar issues arise in diffusion models, where time-conditioning is essential for representing diffusion steps but is not natively accommodated by static LN parameters (Pospisil et al., 2022, Liu et al., 13 Jun 2024).

2. Formulations of TD-LN in RNNs and Generative Models

2.1 Assorted-Time Normalization (ATN): Time-Windowed Statistics

For RNNs, ATN generalizes LN by pooling normalization statistics over a moving window of $k$ recent time steps. For $a^{(t)}$, define a local history

$$A_k^{(t)} = \{a^{(t)}, a^{(t-1)}, \ldots, a^{(t-k+1)}\}$$

Then the ATN mean and variance are:

$$\mu_{t,k} = \frac{1}{nk} \sum_{j=0}^{k-1} \sum_{i=1}^n a_i^{(t-j)}$$

$$\sigma^2_{t,k} = \frac{1}{nk} \sum_{j=0}^{k-1} \sum_{i=1}^n (a_i^{(t-j)} - \mu_{t,k})^2$$

The normalization is then:

$$y^{(t)} = \gamma \odot \frac{a^{(t)} - \mu_{t,k}}{\sqrt{\sigma^2_{t,k}+\epsilon}} + \beta$$

where $\gamma, \beta$ are shared across all time steps. This method enables the post-normalized activations to retain information about slow changes in input statistics over time, breaking time-invariance (Pospisil et al., 2022).

2.2 TD-LN via Time-Parameterized Gain and Bias

In diffusion models, especially where time information must be injected into each normalization layer, TD-LN is realized by making $\gamma, \beta$ explicit functions of the diffusion timestep $t$. Concretely, for $x \in \mathbb{R}^{B\times C \times \dots}$,

$$\text{TD-LN}(x, t) = \gamma(t) \odot \frac{x - \mu(x)}{\sigma(x)} + \beta(t)$$

where $\gamma(t), \beta(t) \in \mathbb{R}^C$ interpolate between two learned endpoints as $t$ varies. The interpolation is governed by a scalar “gate”:

$s(t) = \sigma(w t + b)$

$$\gamma(t) = s(t) p_1 + (1 - s(t)) p_2,\quad \beta(t) = s(t) p_3 + (1 - s(t)) p_4$$

with $p_1, p_2, p_3, p_4 \in \mathbb{R}^C$ and $w,b\in\mathbb{R}$ (Liu et al., 13 Jun 2024).

3. Theoretical Properties

TD-LN methods such as ATN preserve essential invariances of standard LN. Specifically, weight rescaling invariance holds: for a linear preactivation map $H$ scaled by $\delta$, all $a^{(t)}\to \delta a^{(t)}$, which in turn scales $\mu_{t,k}$ and $\sigma_{t,k}$ by $\delta$, leaving $y^{(t)}$ invariant. This property ensures robustness to gradient explosion or vanishing associated with norm fluctuations. Additionally, gradient propagation through ATN remains computationally efficient, requiring only $O(kn)$ overhead per time step. For diffusion TD-LN, the interpolation is directly motivated by PCA analysis showing that 95% of AdaLN’s time-dependent gain/bias variance lies in a two-dimensional subspace, justifying a rank-2 affine parameterization (Pospisil et al., 2022, Liu et al., 13 Jun 2024).

4. Empirical Evaluation and Performance

ATN-LN in RNNs has been evaluated on benchmarks such as the Copying, Adding, Denoise Problems, and language modeling tasks (e.g., PTB and WikiText-2). Across all tested scenarios, ATN-LN strictly outperforms both vanilla LSTM and LN-LSTM, yielding lower minimum train/validation loss or perplexity. Exemplary results for ATN-LN versus LN and LSTM appear in the table:

Task / Metric	LSTM Loss/Perp	LN	ATN
Copying T=100 (CE ×10⁻¹)	1.739/1.731	1.542/1.529	1.354/1.354
Adding T=100 (MSE ×10⁻³)	1.034/1.212	1.319/0.866	0.687/0.385
Denoise T=100 (MSE ×10⁻²)	14.22/14.64	2.489/3.169	1.733/2.073
PTB bpc	1.692/1.743	1.390/1.520	1.381/1.511
WikiText-2 (perplexity)	80.68/65.65	80.24/58.00	78.55/56.06

In diffusion models, full replacement of static LN by TD-LN in both Transformer and ConvNeXt blocks results in state-of-the-art FID scores on ImageNet 256×256 (FID 1.70) and 512×512 (FID 2.89), using significantly fewer parameters than widely-adopted AdaLN schemes. Ablation demonstrates that TD-LN achieves lower FID than AdaLN-Zero while reducing the parameter count by over 60 million per model in tested configurations (Pospisil et al., 2022, Liu et al., 13 Jun 2024).

5. Parameterization Efficiency and Comparison to Prior Time-Conditioning

Conventional AdaLN approaches model $\gamma(t), \beta(t)$ via large MLPs, often incurring $O(C^2)$ or greater parameterization per layer. In contrast, TD-LN parameterizes time-dependence using only $4C+2$ parameters per layer, employing two “endpoints” for both gain and bias and a scalar gating function. TD-LN is therefore substantially more parameter-efficient than AdaLN, offering a strong trade-off between expressivity and architectural simplicity, and avoids the prohibitive memory/compute overheads that destabilize large ConvNeXt-based architectures with O($C^2$) conditioning layers. No elaborate time-token or sinusoidal embedding is necessary: TD-LN treats timestep $t$ as a scalar and applies a minimal affine transformation followed by sigmoid gating (Liu et al., 13 Jun 2024).

6. Practical Implementation Considerations

Deploying TD-LN requires only modest modifications to existing code: in RNNs, ATN incurs $O(nk)$ additional storage and compute per step for holding the temporal history; in diffusion models, each LN layer’s $\gamma, \beta$ are replaced with their time-parameterized counterparts. The per-layer parameters in TD-LN are not shared across layers, and standard optimizers such as AdamW suffice for stable training. Hyperparameter $k$ in ATN is tuned to balance temporal sensitivity against statistical stability, typically $k=10$–60 for sequence modeling. The approach is robust, “out of the box,” and generalizes to both Transformer- and CNN-based architectures (Pospisil et al., 2022, Liu et al., 13 Jun 2024).

7. Summary and Relevance

Time-Dependent Layer Normalization brings temporal awareness to normalization layers in neural network architectures. By either aggregating normalization statistics across time or allowing normalization parameters to vary with time, TD-LN allows learning algorithms to exploit temporally evolving structure in sequential, generative, and diffusion modeling contexts. The architectural modifications are parameter-efficient, preserve key invariances, and yield strict empirical improvements on a variety of tasks compared to static normalization baselines. The methodology is extensible to any statistic-based normalizer and is compatible with online and offline, unidirectional or bidirectional temporal contexts—a flexible and empirically validated normalization paradigm for modern sequential and time-dependent models (Pospisil et al., 2022, Liu et al., 13 Jun 2024).