Adding Dynamic Tanh (DyT) element wise Operation (Paper from Meta)

[coder0143]

Summary

An important paper from Meta(FAIR), MIT, NYU and Princeton authors in late 2025 shows transformers can perform equally if not better using a Dynamic tanh operation instead of a Layer Normalization. The PR for this issue will aim to add this technique to the nnx src (norms) with minimal code. There is a successor paper to this as well featuring Dynamic erf (Derf) function for which i will open another issue.

I am attaching the related links here:
Paper: https://arxiv.org/pdf/2503.10622
Title: Transformers without Normalization
Code: https://github.com/jiachenzhu/DyT
Website: https://jiachenzhu.github.io/DyT/

cc: @cgarciae @vfdev-5 @samanklesaria

Additional context

[Ashutosh0x]

Hi @cgarciae @vfdev-5 @samanklesaria, I'd love to take this on! I plan to implement DyT as an nnx.Module in flax/nnx/nn/normalization.py, following the existing LayerNorm/RMSNorm patterns.

The implementation is straightforward from the paper:

• y = gamma * tanh(alpha * x) + beta
• Learnable parameters: alpha (init 0.5), gamma (init 1.0), beta (init 0.0)
• Supports configurable feature_axes like other norm layers

I'm also an active contributor to google-deepmind/chex (7 merged PRs including JAX 0.8.2 sharding compatibility) so I'm familiar with the ecosystem conventions.

I'll also address the sister issue #5457 (Derf) in the same PR since both are normalization-free pointwise operations with identical API surfaces — this keeps the review efficient and makes the comparison clear.

Will have a PR up shortly with full tests and docstrings. Any design preferences before I start?

[Ashutosh0x]

PR submitted: #5459

Implementation summary:

• DyT as an nnx.Module in flax/nnx/nn/normalization.py
• Follows existing RMSNorm/LayerNorm patterns for parameter init, dtype promotion, and feature_axes broadcasting
• 13 test cases covering initialization, forward pass, gradients, dtypes, JIT, and edge cases
• Exported as nnx.DyT

Design note: alpha is a learnable nnx.Param (not a fixed hyperparameter), matching the paper's approach where each sublayer adapts its own alpha during training. No mask parameter since DyT has no reduction operations.

Happy to iterate on any design feedback, @vfdev-5.

[samanklesaria]

Given that tanh(αx) is such an easy operation to code manually and adding new layers increases the maintenance burden on flax, I'm skeptical that adding such a layer is a good idea right now. No major frontier open models have adopted it (e.g. llama, qwen, mistral); RMSNorm seems the current standard. My impression is that DyT is more sensitive to hyper-parameters at scale, making it less useful in applied settings.

[vfdev-5]

I close this issue as not planned following the same context as here: #5457 (comment)