Normalization के बिना Transformers की तकनीक

• Transformer में लगभग अनिवार्य मानी जाने वाली Layer Norm/RMSNorm की जगह Dynamic Tanh (DyT) लगाने पर भी, मौजूदा normalization मॉडल्स के बराबर या उनसे बेहतर प्रदर्शन मिल सकता है
• DyT DyT(x) = tanh(αx) के रूप में एक element-wise operation है, और यह उस अवलोकन से शुरू होता है कि Transformer के भीतर Layer Normalization अक्सर tanh जैसी S-आकार की input-output mapping बनाता है
• इसका implementation PyTorch की कुछ ही लाइनों में हो जाता है, और trainable alpha, weight, bias के साथ tanh(alpha * x) के output पर scale और bias लागू किया जाता है
• मूल्यांकन ViT, ConvNeXt, MAE, DINO, DiT, LLaMA, wav2vec 2.0, HyenaDNA, Caduceus तक फैला है, यानी vision, language, speech और DNA sequence modeling जैसे व्यापक क्षेत्रों को कवर करता है
• अलग से hyperparameter tuning के बिना भी कई settings में normalization-आधारित समकक्ष मॉडल्स के समान या बेहतर नतीजे मिले, जिससे यह धारणा फिर से सोचने लायक बनती है कि normalization layer अनिवार्य है

---

Dynamic Tanh क्या बदलता है

• DyT एक सरल layer है जो Transformer block में Layer Norm या RMSNorm को replace करती है
• इसका मुख्य operation element-wise लागू होने वाला DyT(x) = tanh(αx) है
• यह दिखाया गया है कि normalization layers हटाने पर भी Transformer, पारंपरिक normalization वाले Transformer जितना या उससे बेहतर प्रदर्शन कर सकता है
• इस विचार की शुरुआत उस अवलोकन से हुई कि Transformer के भीतर Layer Normalization अक्सर scaled tanh function जैसी input-output relationship बनाता है

Implementation का तरीका

• DyT module को PyTorch में छोटे code से implement किया जा सकता है

class DyT(nn.Module):
    def __init__(self, num_features, alpha_init_value=0.5):
        super().__init__()
        self.alpha = nn.Parameter(torch.ones(1) * alpha_init_value)
        self.weight = nn.Parameter(torch.ones(num_features))
        self.bias = nn.Parameter(torch.zeros(num_features))

    def forward(self, x):
        x = torch.tanh(self.alpha * x)
        return x * self.weight + self.bias

• alpha एक trainable parameter है, और इसका initial value 0.5 रखा गया है
• weight और bias भी trainable parameters हैं, जो tanh(alpha * x) के output पर लागू होते हैं

Layer Normalization से निकला अवलोकन

• Transformer का Layer Normalization(LN), scaled tanh function के काफ़ी करीब input-output mapping बनाता है
• शुरुआती layers में यह mapping ज़्यादातर linear के करीब होती है
• जैसे-जैसे layers गहरी होती जाती हैं, tanh function की खास S-आकार की curve और स्पष्ट दिखने लगती है
• इस अवलोकन में Vision Transformer(ViT), speech Transformer मॉडल wav2vec 2.0, और Diffusion Transformer(DiT) की चुनी हुई LN layers शामिल थीं

मूल्यांकन का दायरा और नतीजे

• DyT का मूल्यांकन कई architectures और tasks पर किया गया
  • supervised vision: ViT, ConvNeXt
  • self-supervised vision: MAE, DINO
  • diffusion models: DiT
  • large language models: LLaMA
  • self-supervised speech: wav2vec 2.0
  • DNA sequence modeling: HyenaDNA, Caduceus
• सभी मामलों में DyT लागू किए गए Transformer ने normalization-आधारित समकक्ष मॉडल्स के समान या उनसे बेहतर प्रदर्शन दिखाया
• मूल्यांकन का दायरा recognition से generation, supervised learning से self-supervised learning, और computer vision से language modeling तक व्यापक है

संदर्भ सामग्री

• Download Paper: शोध के पूरे technical details वाला paper
• View on GitHub: implementation details देखने के लिए repository
• View Summary: शोध परिणामों का संक्षिप्त सार
• Transformers without Normalization को CVPR 2025 paper के रूप में स्वीकार किया गया है

[GN⁺]

Hacker News पर राय

• अगर यह सच है, तो यह काफी अच्छा incremental improvement है। लगता नहीं कि इससे मॉडल की performance बहुत meaningful तरीके से बढ़ेगी, लेकिन यह आज के ज़्यादातर cutting-edge LLMs में इस्तेमाल होने वाले RMSNorm से कम compute लेता है, इसलिए training तेज़ और सस्ती हो सकती है
  • हालांकि Transformer की कुल computation में RMSNorm का हिस्सा काफी छोटा होता है। आम तौर पर reduction operations को आगे-पीछे के operations के साथ fuse किया जा सकता है
  • मैंने इसे अभी अपने निजी Transformer training benchmark में लागू करके देखा, और नतीजे बहुत निराशाजनक रहे। RMSNorm इस्तेमाल करने की तुलना में convergence काफी धीमी थी
    alpha को adjust करने से खास असर नहीं पड़ा, इसलिए काफी hyperparameter tuning या ज़्यादा sophisticated initialization की ज़रूरत हो सकती है। PyTorch default initialization और orthogonal initialization, दोनों आज़माए, लेकिन कोई फर्क नहीं पड़ा
    या शायद मेरे इस्तेमाल किए गए scalar optimizer के साथ यह ठीक से match नहीं करता। मैं एक custom scalar optimizer इस्तेमाल करता हूं जो Adam से तेज़ convergence देता है, लेकिन DyT layers में यह Adam जैसा ही दिखा
    या फिर यह ऐसा तरीका हो सकता है जो दसियों अरब tokens के बाद ही catch up करता हो, पर इतना लंबा test करने का budget मेरे पास नहीं है
• float8 जैसे low-precision formats इस्तेमाल करने पर आम तौर पर normalization से पहले activations को BF16 में promote करना पड़ता है। इसलिए जितनी कम precision पर जाते हैं, normalization layer का compute share उतना बढ़ता है
  अगर ऐसी layers को replace किया जा सके, तो compute cost काफी कम करने में मदद मिलेगी
• details पढ़नी पड़ेंगी, लेकिन normalization हटाना बड़ा meaningful हो सकता है। नई architectures आज़माते समय network को ठीक से normalized रखने की tuning हमेशा झंझट भरी होती है
  tanh के भी दूसरे effects होंगे। क्योंकि normalization कभी-कभी conditioning problem को हल कर रहा होता है। फिर भी अधिक alternatives मिलना स्वागतयोग्य है
• तो क्या vanishing gradients अब समस्या नहीं रहे?
  • layers को सही तरह initialize करें, तो deep networks में भी gradient magnitude को vanish या explode होने से बचाए रखा जा सकता है। उदाहरण के लिए, अगर हर layer output का mean 0 और standard deviation 1 पर रखा जाए, तो gradients भी reasonable range में आ जाते हैं
    Kaiming He वगैरह का मूल ResNet paper और उसके बाद के papers recommend करूंगा
    RNNs के modern approach के लिए DeepMind का https://arxiv.org/abs/2303.06349 पढ़ने लायक है
    मुख्य बात यह है कि सबसे बड़ा eigenvalue, यानी spectral radius, 1 के आसपास होना चाहिए। इसका मतलब है कि linear transformation को बार-बार apply करने पर activations न बढ़ें, न घटें
  • ResNet ने vanishing gradients को लगभग हल कर दिया माना जा सकता है। exploding gradients को आम तौर पर अच्छे parameter initialization और normalization से handle किया जाता है। यह paper असल में normalization का alternative propose कर रहा है
  • अच्छा सवाल है। वह समस्या उस दौर की थी जब tanh को activation function के रूप में इस्तेमाल किया जाता था, और residual connections व normalization layers से पहले की बात थी। दूसरी activation functions और residual connections के साथ tanh को normalization की तरह इस्तेमाल करना ठीक लगता है
  • diagram में दिखता है कि Transformer residual सीखता है। इसका रूप y = x + f(x) है
• पता नहीं सिर्फ मुझे ऐसा लग रहा है या नहीं, लेकिन paper के graph में LNinput और LNoutput की तुलना करते हुए tanh(a*x) के बाद भी weight और bias लगाए गए लगते हैं
  similarity देखने के लिए LayerNorm output से weight और bias हटाए गए result से तुलना करनी चाहिए, ऐसा नहीं है क्या?
  final results अच्छे हों तो फर्क नहीं पड़ता, लेकिन असल में बदले जा रहे हिस्से को अलग से देखें तो बेहतर समझ आ सकती है कि क्या हो रहा है
  • implementation देखने पर लगता है कि tanh compute करने के बाद weight और bias apply किए जाते हैं
• practical तौर पर इसका मतलब क्या है?
  • abstract के मुताबिक, DyT लगाने पर normalization-free Transformer, ज़्यादातर hyperparameter tuning के बिना भी normalized counterpart models के बराबर या उनसे बेहतर performance दे सकता है