Tag: Deep Learning

  • Differential Transformer: Canceling Context Noise with Differential Attention

    Differential Transformer: Canceling Context Noise with Differential Attention

    TL;DR

    • The Differential Transformer cancels attention noise and focuses attention by calculating scores as the difference between two separate softmax attention maps.
    • It achieves standard Transformer performance with approximately 35% fewer parameters or training tokens while reducing contextual hallucinations and enhancing long-context retrieval.
    • The DEX adaptation method enables existing pretrained language models to effectively incorporate key advantages of differential attention using minimal adaptation data (< 0.01%).

    Introduction to the Differential Transformer

    Developed by researchers from Microsoft Research and Tsinghua University—including Tianzhu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei—the Differential Transformer (Diff Transformer) introduces a foundational architectural modification to standard Transformer models. Accepted as an Oral Presentation at ICLR 2025, the architecture addresses an inherent limitation in standard softmax attention: attention heads tend to overallocate probability weight to irrelevant context, which is often viewed as a consequence of softmax normalizing scores into strictly positive values summing to 1. In the authors’ motivation, this attention noise may relate to observed challenges such as contextual hallucinations, prompt-order sensitivity, and the “lost-in-the-middle” phenomenon.

    Drawing inspiration from noise-canceling headphones and differential amplifiers in electrical engineering, the Differential Transformer incorporates a Differential Attention Mechanism designed to cancel common-mode background noise and extraneous context. The architecture demonstrates several key advantages across language modeling tasks:

    • Improved Scaling Efficiency: The Diff Transformer matches the performance of standard Transformers while requiring approximately 35% fewer parameters or training tokens.
    • Long-Context Robustness: The model demonstrates stronger retrieval performance on needle-in-a-haystack tasks and long-context evaluations up to 64k tokens, showing less degradation across needle depths and context lengths.
    • Reduced Contextual Hallucinations: By being less distracted by irrelevant context, it decreases hallucination rates on text summarization and question-answering tasks.
    • In-Context Learning Stability: The architecture mitigates sensitivity to the permutation order of few-shot examples, offering higher stability during prompt reordering.
    • Quantization Readiness: Alongside noise cancellation, the architecture suppresses extreme activation outliers, facilitating cleaner post-training low-bit quantization as a distinct representational benefit.

    Sources: [2410.05258] Differential Transformer, Differential Transformer, microsoft.com/en-us/research/publication/differential-transformer, Microsoft’s Differential Transformer cancels attention noise in LLMs, OpenReview (paper discussion page)

    Cancelling Context Noise Through Attention Map Subtraction

    At the heart of these performance gains is the specific way Differential Attention operates. The Differential Attention Mechanism cancels attention noise by calculating attention scores as the difference between two separate softmax attention maps:

    DiffAttn(X) = (softmax(Q1 * K1^T / sqrt(d)) - lambda * softmax(Q2 * K2^T / sqrt(d))) * V

    where query and key projections are split into two distinct paths (Q1, K1 and Q2, K2) and λ is a learnable scalar parameter. In practice, λ is reparameterized via learnable vectors λ_q1, λ_k1, λ_q2, and λ_k2 (specifically computed as λ = exp(λ_q1 · λ_k1) – exp(λ_q2 · λ_k2) + λ_init, where λ_init is set according to the layer index). This subtraction directly cancels common-mode noise and counteracts the tendency of conventional architectures to overallocate attention to irrelevant context, amplifying focus on relevant tokens and promoting the emergence of sparse attention patterns.

    By canceling attention noise through this dual-map subtraction, the Diff Transformer achieves consistent empirical advantages across both scaling and downstream evaluations, outperforming standard Transformer baselines across various model sizes and training token volumes.

    Sources: Differential Transformer

    Follow-Up Analysis and Adapting Pretrained Models

    While the original Differential Transformer paper framed its mechanism primarily around noise cancellation, an independent follow-up study titled “Understanding Differential Transformer Unchains Pretrained Self-Attentions” (arXiv:2505.16333) by Chaerin Kong, Jiho Jang, and Nojun Kwak investigated the underlying factors behind its empirical success. The authors uncovered three key factors:

    • Enhanced expressivity via negative attention
    • Reduced redundancy among attention heads
    • Improved learning dynamics

    Because the standard Differential Transformer architecture requires training from scratch, which limits its direct use with existing pretrained models, the study proposed DEX, an efficient adaptation method designed to integrate the advantages of differential attention into pretrained models.

    Rather than computing two separate softmax maps, DEX reuses the original softmax attention scores and adds a lightweight differential operation on the output value matrix. This formulation allows pretrained language models to effectively incorporate key advantages of differential attention while remaining lightweight in training and inference, achieving significant performance gains across diverse benchmarks with minimal adaptation data (< 0.01%).

    Sources: Understanding Differential Transformer Unchains Pretrained Self-Attentions

    Closing thoughts

    In conclusion, by subtracting attention noise through two softmax maps, the Differential Transformer addresses a key limitation of standard softmax attention. Its empirical gains across scaling efficiency, long-context retrieval, and contextual hallucination reduction, alongside activation outlier suppression, highlight the advantages of differential attention representations. Furthermore, independent analysis reveals that benefits such as negative attention expressivity, reduced head redundancy, and improved learning dynamics also underpin its performance, inspiring practical adaptation methods like DEX for pretrained models. Differential attention thus stands as a grounded architectural refinement that delivers tangible, parameter- and data-efficient performance gains.

    Frequently Asked Questions

    What is the primary problem the Differential Transformer solves?

    The Differential Transformer addresses attention noise in standard softmax attention, where attention heads tend to overallocate probability weight to irrelevant context, often seen as a consequence of softmax normalizing scores into strictly positive values summing to 1.

    How does the Differential Attention Mechanism cancel context noise?

    It computes query and key projections across two separate paths, takes the difference between two distinct softmax attention maps weighted by a learnable parameter λ, and subtracts common-mode background noise to focus on salient tokens.

    What performance advantages does the Differential Transformer provide?

    It matches standard Transformer performance using approximately 35% fewer parameters or training tokens, improves long-context retrieval across various depths and lengths up to 64k tokens, lowers contextual hallucination rates in QA and summarization, stabilizes in-context learning against prompt reordering, and suppresses activation outliers for cleaner low-bit quantization.

    What is the DEX method and why is it needed?

    Because the original Differential Transformer architecture requires training from scratch, the DEX method was proposed in an independent follow-up study (arXiv:2505.16333 by Chaerin Kong, Jiho Jang, and Nojun Kwak) to adapt pretrained language models. Instead of computing two softmax maps, DEX reuses existing softmax attention scores and applies a lightweight differential operation on the output value matrix using minimal adaptation data (< 0.01%).


    Keep reading: Training & Architectures

    All guides in this series: LLM Training & Architecture Guides: Fine-Tuning, RL & New Models

  • Muon Optimizer: Scaling Matrix Orthogonalization for Efficient Deep Learning Training

    Muon Optimizer: Scaling Matrix Orthogonalization for Efficient Deep Learning Training

    TL;DR

    • The Muon optimizer trains 2D hidden layer weight matrices via steepest descent under the spectral norm using efficient iterative Newton-Schulz orthogonalization.
    • Muon operates as a hybrid optimizer: it manages 2D hidden-layer weight matrices, while AdamW handles everything else — genuinely 1D parameters (biases, normalization scales) plus the input embedding and output layers, which are excluded from Muon for role-based reasons even though they are typically 2D matrices themselves.
    • Compared to AdamW, Muon cuts training compute — about 1.35x faster wall-clock in Keller Jordan’s NanoGPT speedrun and up to 2x fewer FLOPs under Moonlight’s compute-optimal scaling law — while also cutting VRAM overhead by eliminating the second-moment buffer and reducing the high condition number (near low-rank concentration) of raw weight updates.

    Introduction: Understanding the Muon Optimizer and Orthogonal Updates

    The Muon Optimizer—short for MomentUm Orthogonalized by Newton-Schulz—is an optimization method designed for training the Hidden Layers of neural networks. Unlike conventional first-order optimizers such as SGD, Adam, and AdamW, which treat weight parameters as flattened 1D Parameters and apply coordinate-wise scaling, Muon treats parameters as 2D matrices. It performs Steepest Descent with respect to the Spectral Norm (the matrix operator norm), optimizing both the directional geometry and the magnitude of Weight Updates.

    Muon was introduced in late 2024 by Keller Jordan and collaborators, building on foundational theoretical work by Jeremy Bernstein and Laker Newhouse (‘Old Optimizer, New Norm: An Anthology’), which frames this Steepest Descent step under an RMS-to-RMS operator norm — a normalized formulation that is effectively equivalent to the classical Spectral Norm used elsewhere in this article, differing mainly in how the norm is scaled by matrix dimensions rather than in the underlying geometry. It was first benchmarked at scale in the open-source Modded-NanoGPT speedrun repository. The optimizer has since been expanded through theoretical analyses, such as the Newton–Muon optimizer studied by Du and Su at UPenn, and scaled by Moonshot AI (the Kimi Team) to train the Moonlight 3B/16B Mixture-of-Experts (MoE) model across 5.7 trillion tokens.

    At its core, Muon operates through a multi-step update pipeline:

    • Momentum Collection: Gradient updates are first accumulated using standard Nesterov momentum.
    • Newton-Schulz Orthogonalization: Rather than computing an exact, computationally expensive Singular Value Decomposition (SVD) on GPUs, Muon applies an iterative Newton-Schulz Iteration (typically 5 steps). This quickly drives all singular values of the update matrix close to 1.0, generating an approximate Matrix Orthogonalization equivalent to the matrix-sign function and Polar Decomposition.
    • Aspect-Ratio Scaling: The orthogonalized update is rescaled according to matrix dimensions to preserve stable signal propagation across layers with varying shapes. Keller Jordan’s original implementation uses a ratio-based factor, sqrt(max(1, rows/cols)); the Moonlight paper (Kimi Team) proposes an explicit coefficient form instead, 0.2 * sqrt(max(A, B)), noting the two scalings are equivalent up to a global scale only when the matrices share the same second dimension, and can otherwise diverge.

    Muon is implemented as a hybrid optimizer. It is specialized strictly for 2D weight matrices in Hidden Layers and linear layers—such as Multi-Layer Perceptron (MLP) weights, attention projections in Transformer Architectures, and reshaped convolutional kernels. Everything else is optimized concurrently using a traditional optimizer like AdamW: genuinely 1D parameters such as Biases and Normalization Scales, plus the input embedding and output (unembedding) layers. Those embedding and output layers are typically 2D matrices themselves, but Muon’s authors exclude them for role-based reasons rather than dimensionality—embeddings on the theoretical grounds of the modular norm framework, and output layers on purely empirical grounds.

    This Matrix Orthogonalization scheme provides several key advantages over standard coordinate-wise optimizers:

    • Compute Efficiency and Convergence: Reported efficiency gains come from two distinct sources — Keller Jordan’s NanoGPT speedrun recorded roughly 1.35x faster wall-clock training time, while the Moonlight paper’s compute-optimal scaling law separately found up to 2x fewer training FLOPs versus AdamW to reach the same validation loss.
    • Lower Optimizer Memory Overhead: Muon tracks only a single first-moment momentum buffer per parameter and eliminates the second-moment buffer (v_t) required by AdamW, yielding notable VRAM savings on Hidden Layers.
    • Reducing Update Condition Number: By normalizing singular values, the optimizer prevents Weight Updates from collapsing into a small subset of high-variance singular directions—a near low-rank, high-condition-number pathology of raw gradients—better engaging the full representational capacity of weight matrices.
    • Empirical Performance: In addition to setting speedrun training records on single-node setups for CIFAR-10 and GPT-2 (small and medium), Muon has demonstrated stable scaling in multi-billion-parameter LLM pretraining. The Moonlight paper’s Distributed Muon implementation follows ZeRO-1 to partition optimizer states across the data-parallel dimension — with tensor parallelism usable only as a complementary, non-essential option — combined with adjusted weight decay.

    Sources: Muon: An optimizer for hidden layers in neural networks, GitHub – KellerJordan/Muon: Muon is an optimizer for hidden layers in neural networks · GitHub, Deriving Muon, medium.com/@md.abir1203/muon-optimizer-the-complete-guide-to-2-faster…, [2502.16982] Muon is Scalable for LLM Training

    Closing thoughts

    Taken together, these technical strengths illustrate how shifting from coordinate-wise updates to geometry-aware 2D matrix orthogonalization enables Muon to directly address the structural properties of hidden layers, yielding clear optimization advantages over standard first-order methods. Its hybrid architecture—applying efficient Newton-Schulz iterations to weight matrices while retaining AdamW for 1D parameters and embeddings—strikes a sensible balance between theoretical rigor and implementation practicality.

    Ultimately, with demonstrated FLOP reductions and lower memory overhead validated from single-node speedruns up to multi-trillion-token LLM pretraining, Muon establishes matrix-level spectral descent as a genuinely viable, high-efficiency paradigm for training modern neural networks.

    Frequently Asked Questions

    What is the Muon optimizer?

    Muon (MomentUm Orthogonalized by Newton-Schulz) is an optimization method designed for training hidden layers of neural networks by treating parameters as 2D matrices rather than flattened 1D vectors. It performs steepest descent with respect to the spectral norm to optimize the directional geometry and magnitude of weight updates.

    How does Muon achieve matrix orthogonalization efficiently?

    Instead of computing an expensive Singular Value Decomposition (SVD) on GPUs, Muon applies an iterative Newton-Schulz iteration, typically in 5 steps. This quickly drives update matrix singular values close to 1.0, providing an approximate orthogonalization equivalent to the matrix-sign function and polar decomposition.

    Why is Muon used as a hybrid optimizer alongside AdamW?

    Muon is specialized strictly for 2D weight matrices in hidden layers, such as MLP weights and attention projections. Everything else uses a traditional optimizer like AdamW instead: 1D parameters such as biases and normalization scales, plus the input and output embedding layers. Those embedding layers are typically 2D themselves, but are excluded from Muon for role-based reasons—not because they are non-2D.

    What are the primary benefits of Muon compared to AdamW?

    Muon reduces training compute versus AdamW — about 1.35x faster wall-clock time in Keller Jordan’s NanoGPT speedrun, and up to 2x fewer FLOPs under Moonlight’s compute-optimal scaling law — while also lowering VRAM usage by eliminating AdamW’s second-moment buffer. Additionally, it helps reduce the high condition number of raw weight updates by preventing them from collapsing into a small subset of high-variance singular directions.


    Keep reading: Training & Architectures

    All guides in this series: LLM Training & Architecture Guides: Fine-Tuning, RL & New Models

  • Fixing Silent PyTorch DataLoader Crashes in Docker Containers

    Fixing Silent PyTorch DataLoader Crashes in Docker Containers

    TL;DR

    • PyTorch DataLoader uses POSIX shared memory at /dev/shm to efficiently pass CPU tensor data between parallel worker processes.
    • Docker containers allocate a default limit of only 64 MB to /dev/shm, causing workers to crash with a SIGBUS signal when mapped pages are accessed without backing space — a failure mode distinguishable in logs from a SIGKILL caused by host or container RAM exhaustion.
    • Preventing these DataLoader crashes requires configuring container-level shared memory to match PyTorch’s multi-process requirements.

    Understanding the Cause and Symptoms of DataLoader Shared Memory Crashes

    PyTorch supports multi-process data loading through torch.utils.data.DataLoader when the num_workers parameter is set to a value greater than zero. Under this configuration, parallel worker processes fetch and collate data simultaneously. To efficiently pass CPU tensor data between worker processes and the main process without expensive data copying, PyTorch relies on Inter-Process Communication (IPC) backed by POSIX shared memory, which is mounted at /dev/shm.

    By default, Docker allocates a restricted size of only 64 MB to /dev/shm for Docker containers. When deep learning workloads involve large batch sizes, heavy tensor payloads, or multiple worker processes, the active data buffers can quickly exceed this 64 MB shared memory allocation. In tmpfs-backed shared memory, mapping operations such as mmap or ftruncate may initially succeed, but when worker processes touch and access mapped pages without sufficient backing space, the operating system raises a SIGBUS (Bus Error) signal. This terminates worker processes with error notices such as DataLoader worker (pid X) is killed by signal: Bus error. It is possible that dataloader's workers are out of shared memory… or DataLoader worker exited unexpectedly. This failure is distinct from a SIGKILL signal (killed by signal: Killed), which indicates that an out-of-memory killer terminated a process because actual container RAM — a separate resource from /dev/shm — was exhausted.

    Resolving Shared Memory Exhaustion

    To prevent these crashes, you can configure the container’s shared memory capacity or IPC settings at launch:

    1. Increase Container Shared Memory Size (Recommended)
      Explicitly set /dev/shm when starting the container using --shm-size:
    docker run --shm-size=8g my_pytorch_image

    Note that 8g is only an illustrative example; you should size this value according to your number of workers, batch size, and prefetch settings. Prefer this method in most production workflows because it preserves container isolation while providing dedicated shared memory.

    In docker-compose.yml, specify shm_size for the corresponding service:

    services:
  training:
    image: my_pytorch_image
    shm_size: '8gb'
    1. Use the Host IPC Namespace
      Alternatively, share the host machine’s IPC namespace directly with the container:
    docker run --ipc=host my_pytorch_image

    Prefer --ipc=host in trusted environments where you want to grant the container access to host IPC resources without manually sizing /dev/shm, keeping in mind that sharing the host IPC namespace weakens container isolation.

    Sources: docs.pytorch.org/docs/stable/data.html, docs.pytorch.org/docs/stable/notes/multiprocessing.html, Shared Memory in Docker with the shm-size Flag – datawookie, Docker –shm-size: Set /dev/shm and Fix Crashes, PyTorch: increasing –shm-size to allow multiprocessing data loaders · Issue #937 · aws/sagemaker-python-sdk ·

    Closing thoughts

    Ultimately, these crashes highlight a mismatch between PyTorch’s multi-process IPC design and Docker’s restrictive default container settings. The default 64 MB shared memory cap can be problematic when using multiple workers, as signals like SIGBUS or unexpected worker exit messages can lead developers to investigate their dataset code rather than the container configuration. By adjusting container shared memory via --shm-size or using --ipc=host in trusted setups, teams can ensure reliable multi-process data loading without running out of shared memory.

    Frequently Asked Questions

    Why does PyTorch DataLoader crash in default Docker containers?

    PyTorch uses POSIX shared memory at /dev/shm to pass tensor data between worker processes and the main process when num_workers is greater than zero. Docker restricts /dev/shm to 64 MB by default, which is easily exhausted during heavy data loading workloads, leading to worker crashes.

    What default shared memory allocation does Docker provide to containers?

    By default, Docker allocates a restricted size of only 64 MB to /dev/shm.

    What error signals or messages occur when PyTorch shared memory is exhausted?

    Worker processes terminate with a SIGBUS (Bus Error) signal, often accompanied by messages stating that the worker was killed by a bus error and may be out of shared memory, or generic messages such as “DataLoader worker exited unexpectedly”. A SIGKILL signal is a distinct failure mode caused by out-of-memory termination due to actual RAM exhaustion, not /dev/shm size.

    How can I fix /dev/shm exhaustion in Docker?

    You can override the default shared memory limit by passing --shm-size to docker run (sized according to your number of workers, batch size, and prefetch), setting shm_size in docker-compose.yml, or using --ipc=host to share the host IPC namespace in trusted environments.

    Why does PyTorch DataLoader rely on POSIX shared memory?

    PyTorch relies on Inter-Process Communication (IPC) backed by POSIX shared memory to pass CPU tensor data between worker processes and the main process. This avoids expensive data copying when running multi-process data loading.


    Keep reading: Training & Architectures

    All guides in this series: LLM Training & Architecture Guides: Fine-Tuning, RL & New Models