DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention

[remyx-ai]

Model/Pipeline/Scheduler description

DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention (arXiv:2405.18428) swaps a Diffusion Transformer block's softmax self-attention for a sub-quadratic Gated Linear Attention (GLA) recurrence. Proposal: add a parameter-free GatedLinearAttnProcessor2_0 to src/diffusers/models/attention_processor.py.

• A per-head recurrent state of shape (head_dim, head_dim) replaces softmax's O(T²) score matrix, so compute is O(T · head_dim²) and no T × T matrix is ever materialised.
• A data-dependent forget gate modulates the state — the GLA ingredient that lets linear attention approach softmax quality.
• Drop-in for AttnProcessor2_0: reuses the Attention module's to_q/to_k/to_v/to_out projections and adds no parameters, so model.set_attn_processor(GatedLinearAttnProcessor2_0()) is a swap that does not change the parameter count.
• Naive pure-torch recurrence for correctness; the constant-factor speedup from a chunked / associative-scan kernel (e.g. flash-linear-attention) is intentionally deferred.

Proposed implementation (working draft, 207-line diff)

Apply locally with git apply after saving the block below.

diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py
index e2ece5c..73cbc3e 100755
--- a/src/diffusers/models/attention_processor.py
+++ b/src/diffusers/models/attention_processor.py
@@ -5390,6 +5390,119 @@ class SanaLinearAttnProcessor2_0:
         return hidden_states
 
 
+def gated_linear_attention(query, key, value, eps=1e-6):
+    r"""
+    Gated Linear Attention recurrence — DiG (arXiv:2405.18428).
+
+    Sub-quadratic in the sequence length: a per-head recurrent state of shape
+    (head_dim, head_dim) replaces softmax's O(T^2) score matrix, so compute is
+    O(T · head_dim^2) and no T × T matrix is ever materialised. A data-dependent
+    forget gate modulates the state — the GLA ingredient that lets linear
+    attention approach softmax quality. Inputs/outputs are (batch, heads, seq, head_dim).
+    """
+    # Non-negative feature map (linear-attention convention; same role as F.relu
+    # in SanaLinearAttnProcessor2_0). elu(.) + 1 keeps features in [0, +inf).
+    query_feat = F.elu(query) + 1.0
+    key_feat = F.elu(key) + 1.0
+    # Data-dependent forget gate on the recurrent state. The learned gate
+    # projection used when training a DiG checkpoint is omitted here so the
+    # processor stays a parameter-free drop-in for AttnProcessor2_0.
+    gate = torch.sigmoid(query * key)
+
+    orig_dtype = query_feat.dtype
+    query_feat, key_feat, value, gate = (t.float() for t in (query_feat, key_feat, value, gate))
+    lead, seq_len, head_dim = query_feat.shape[:-2], query_feat.shape[-2], query_feat.shape[-1]
+    q = query_feat.reshape(-1, seq_len, head_dim)
+    k, v, g = (t.reshape(-1, seq_len, head_dim) for t in (key_feat, value, gate))
+    n = q.shape[0]
+    state = torch.zeros(n, head_dim, head_dim, device=q.device, dtype=torch.float32)
+    denom = torch.zeros(n, head_dim, device=q.device, dtype=torch.float32)
+    outputs = []
+    for t in range(seq_len):
+        q_t, k_t, v_t, g_t = q[:, t], k[:, t], v[:, t], g[:, t]
+        # forget + write: state <- gate ⊙ state + k ⊗ v
+        state = g_t[:, :, None] * state + k_t[:, :, None] * v_t[:, None, :]
+        # matching gated recurrence for the per-query normaliser denominator
+        denom = g_t * denom + k_t
+        num = torch.bmm(q_t.unsqueeze(1), state).squeeze(1)
+        outputs.append(num / (q_t * denom).sum(-1, keepdim=True).clamp(min=eps))
+    return torch.stack(outputs, dim=1).reshape(*lead, seq_len, head_dim).to(orig_dtype)
+
+
+class GatedLinearAttnProcessor2_0:
+    r"""
+    Processor for implementing Gated Linear Attention (GLA) self-attention.
+
+    Adapted from DiG (arXiv:2405.18428): swaps a Diffusion Transformer block's
+    softmax self-attention for a sub-quadratic gated linear-attention recurrence.
+    Drop-in for `AttnProcessor2_0` — reuses the `Attention` module's
+    `to_q`/`to_k`/`to_v`/`to_out` projections and adds no parameters, so
+    `model.set_attn_processor(GatedLinearAttnProcessor2_0())` is a swap that does
+    not change the parameter count. Naive pure-torch recurrence for correctness;
+    the constant-factor speedup from a chunked / associative-scan kernel (e.g.
+    flash-linear-attention) is intentionally deferred.
+    """
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        temb: torch.Tensor | None = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        batch_size = hidden_states.shape[0]
+        query = attn.to_q(hidden_states)
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        hidden_states = gated_linear_attention(query, key, value)
+        hidden_states = hidden_states.to(query.dtype)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
 class PAGCFGSanaLinearAttnProcessor2_0:
     r"""
     Processor for implementing scaled dot-product linear attention.
@@ -5661,6 +5774,7 @@ AttentionProcessor = (
     | SlicedAttnProcessor
     | SlicedAttnAddedKVProcessor
     | SanaLinearAttnProcessor2_0
+    | GatedLinearAttnProcessor2_0
     | PAGCFGSanaLinearAttnProcessor2_0
     | PAGIdentitySanaLinearAttnProcessor2_0
     | SanaMultiscaleLinearAttention
diff --git a/tests/models/transformers/test_models_dit_transformer2d.py b/tests/models/transformers/test_models_dit_transformer2d.py
index 473a876..830528a 100644
--- a/tests/models/transformers/test_models_dit_transformer2d.py
+++ b/tests/models/transformers/test_models_dit_transformer2d.py
@@ -18,6 +18,7 @@ import unittest
 import torch
 
 from diffusers import DiTTransformer2DModel, Transformer2DModel
+from diffusers.models.attention_processor import Attention, GatedLinearAttnProcessor2_0
 
 from ...testing_utils import (
     enable_full_determinism,
@@ -91,6 +92,62 @@ class DiTTransformer2DModelTests(ModelTesterMixin, unittest.TestCase):
     def test_effective_gradient_checkpointing(self):
         super().test_effective_gradient_checkpointing(loss_tolerance=1e-4)
 
+    def _build_model_and_inputs(self, **init_overrides):
+        init_dict, inp

Open source status

• The model implementation is available. — ⚠️ but no LICENSE file detected on the upstream repo (compat 0.00, class missing); treat as blocking for redistribution/modification until upstream adds a license.
• The model weights are available (Only relevant if addition is not a scheduler).

Provide useful links for the implementation

• Paper: DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention (arXiv:2405.18428)
• Reference implementation: https://github.com/hustvl/DiG
• License: 🔴 none detected upstream — see the caveat under Open source status.

Why this is an Issue, not a PR

The missing upstream LICENSE (no legal permission to redistribute or modify the reference code) is blocking, so the working draft above is shared here for discussion rather than opened as a pull request. Reopen this Issue if you'd like the implementation revisited once upstream licensing is resolved.

Drafted by Outrider — paper: arXiv:2405.18428.

Discovery context

• Confidence: moderate (Remyx relevance 0.65)
• Research interest: pin-arxiv
• Surfaced by Outrider deep-search refine query pin-arxiv:2405.18428 against /search/assets. The engine's normal ranking did not place this paper in the interest's broad pool — it's here because the audit pass identified an under-represented theme this paper covers.