PaperlayerView exact revision on arXiv

Source-linked AI summary

RoLA: Rotary-Positioned Low-Rank Linear Attention for Efficient Diffusion Transformers

Zekun Zhang, Yixiang Cai, Yuxi Liu, Tengxu Sun, Tianle Liu, Zhoutong Wu, Haoyu Li, Baole Ai, Ang Wang, Jiamang Wang, Lin Qu, Kun Yuan

arXiv:2609.06712v3cs.CV

TL;DR

Dense video-DiT attention has quadratic sequence-length cost, and sparse–linear hybrids face difficulty preserving both reusable aggregation and 3D relative rotary geometry. RoLA applies rank-truncated RoPE after the low-rank nonlinearity to retain a reusable global branch without auxiliary positional parameters. Evaluations report favorable quality–efficiency trade-offs under high sparsity, while the evaluation scope is two Wan models.

  • Problem

    Dense attention has O(L^2) cost, and nonlinear feature maps generally do not commute with RoPE, complicating preservation of relative geometry in a reusable linear summary.

    Source: §3 · p. 4

  • Method

    RoLA applies rank-truncated 3D RoPE outside the nonlinear low-rank activation to retain a reusable linear-time global summary without auxiliary absolute-position parameters.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    Mechanistic and generation evaluations report favorable quality–efficiency trade-offs under high sparsity, with generation fidelity retained while computational overhead is substantially reduced.

    Source: §1 · p. 2, §6 · p. 11

  • Takeaways & Limitations

    RoLA provides a sparse-video-DiT construction that keeps genuine cross-token global aggregation compatible with a pretrained 3D-RoPE backbone.

    Source: §1 · p. 2

    The current evaluation focuses on two Wan models.

    Source: §6 · p. 11

Abstract

from arXiv · show

Diffusion Transformers (DiTs) achieve strong video generation quality, but their dense spatiotemporal self-attention scales quadratically with sequence length and quickly becomes the dominant inference bottleneck. Sparse low-rank hybrids alleviate this cost by combining a local sparse branch with a global compressed branch. In video DiTs equipped with 3D Rotary Position Embeddings (RoPE), the global branch faces a structural compatibility issue: when RoPE is applied before a nonlinear feature map, the rotation and nonlinearity generally do not commute, making it difficult to keep a query-independent linear summary while preserving relative rotary geometry. Existing work often sidesteps this issue by replacing genuine cross-token global aggregation with coordinate-conditioned surrogates or learnable absolute positional modules. These compromises can be effective, but they approximate relative decay from absolute coordinates and introduce extra positional parameters. We propose \textbf{RoLA}, a rotary-positioned low-rank linear-attention branch that keeps genuine cross-token aggregation while remaining compatible with a reusable linear summary. The design applies RoPE \emph{outside} the nonlinear low-rank feature map and reuses a truncated subset of the pre-trained rotary schedule matched to the low-rank bottleneck. This yields a linear-time low-rank global branch with relative positional behavior by design and no additional positional parameters; the full sparse--low-rank module still includes the fixed-sparsity sparse branch. Experiments on open-source video DiTs show that the resulting method remains competitive in generation quality at 90\% sparsity while achieving 2.63$\times$ end-to-end inference speedup on Wan2.1-14B (720p, 81 frames, measured on an NVIDIA H100 GPU).

1 Introduction

Dense video-DiT attention is costly, while aggressive sparsification can discard long-range context and conventional sparse–linear hybrids face a conflict between reusable aggregation and 3D relative-position geometry. RoLA applies RoPE after the low-rank feature map to retain a reusable global branch with relative-position structure, without additional positional parameters.

1 Introduction

  • Dense spatiotemporal self-attention scales quadratically with sequence length, while aggressive sparsification can remove interactions important for semantic consistency, motion coherence, and stable scene layout.Sparse attention reduces cost by computing only selected query–key interactions, but the cited concern is at aggressive sparsity.Source: §1 · p. 2
  • In 3D-RoPE video DiTs, position-dependent rotations generally do not commute with nonlinear feature maps, complicating a query-independent global summary that preserves relative geometry.The paper calls this tension the RoPE Dilemma.Source: §1 · p. 2
  • Coordinate-conditioned surrogates replace genuine cross-token aggregation and approximate relative decay from absolute coordinates, while adding positional parameters and giving up reusable linear-attention global context.The passage describes this as one practical workaround.Source: §1 · p. 2
  • RoLA applies RoPE after the nonlinear low-rank feature map and reuses a rank-truncated subset of the pretrained rotary schedule, preserving relative-position structure and a reusable linear-time global summary without new positional parameters.The design is tailored to video DiTs and a sparse deployment engine, rather than claimed as a generic linear relative-position encoder.Source: §1 · p. 2
  • Mechanistic tests on real features and generation experiments across representative backbones report favorable quality–efficiency trade-offs under high sparsity and full-stack inference.The introduction describes this as empirical validation of the method.Source: §1 · p. 2

2 Related Work

Related work spans sparse-only video attention, sparse–linear compensation, and relative-position methods for linear attention. The section distinguishes RoLA’s genuine cross-token global routing from approaches that use coordinate-conditioned or absolute-position surrogates.

Sparse attention for video DiTs

  • Sparse-only video-DiT methods exploit structured, learned, or dynamic sparsity, but under aggressive sparsity may discard weak yet collectively important long-range context.Examples include spatial–temporal sparsity, attention tiles, sparse patterns, block mixtures, distance-aware computation, and dynamic sparsity prediction.Source: §2 · p. 3

Sparse linear hybrids

  • Sparse–linear hybrids compensate for discarded global interactions, but conventional nonlinear feature maps can make preserving 3D-RoPE geometry in a reusable linear branch difficult.RoPeSLR [14] instead uses a per-token coordinate MLP and learnable absolute 3D positional embedding, while RoLA retains cross-token linear aggregation.Source: §2 · p. 3

Rotary position embeddings

  • LRPE [17] preserves relative positional modeling with linear-complexity computation, whereas RoLA is presented as a sparse-video-DiT construction compatible with a pretrained 3D-RoPE backbone and sparse engine.The passage says the per-axis frequency budget is small because 3D RoPE splits the head dimension across temporal and spatial axes.Source: §2 · p. 3

Deployment efficiency

  • Figure 1 organizes deployment-oriented efficiency views into fixed-setting latency, theoretical self-attention FLOPs per denoising step, and long-sequence latency.Its RTX 5090 setup is identified as different from the benchmark evaluation setting.Source: Figure 1 · p. 3

Method overview

  • Figure 2 depicts dense 3D-RoPE attention as sparse top-k spikes plus a residual background, then shows sparse–low-rank branches with gated fusion reusing backbone RoPE without additional positional embeddings.The right side presents the rotary low-rank branch and distribution-aware gated fusion.Source: Figure 2 · p. 4

3 Preliminary

The preliminary section frames full attention as a quadratic-cost operation over flattened spatiotemporal video tokens and explains how 3D RoPE encodes relative position. It then identifies the incompatibility between nonlinear feature maps and reusable linear summaries that motivates RoLA.

Full attention

  • A video is flattened into L spatiotemporal tokens, and full attention has O(L^2) cost that becomes a bottleneck for ultra-long, high-resolution sequences.The notation defines L as the product of temporal and spatial token dimensions.Source: §3 · p. 4

3D RoPE

  • 3D RoPE partitions each attention head across temporal and spatial axes and uses orthogonal rotations to express attention scores through relative position offsets.The relative-position identity is given in equation (1).Source: §3 · p. 4

RoPE Dilemma

  • A pointwise nonlinear feature map generally does not commute with RoPE, making a reusable query-independent summary difficult to reconcile with relative rotary geometry.The paper motivates applying RoPE after the low-rank nonlinearity to inject position information without destroying the reusable summary.Source: §3 · p. 4

Positional kernels

  • With content-controlled inputs, the outside-activation RoPE low-rank map closely matches full attention’s banded relative-distance pattern; inside-activation RoPE distorts it, and removing RoPE yields a nearly flat response.The figure comparison concerns spatial pattern rather than absolute color ranges.Source: Figure 3 · p. 5

Efficient-attention comparison

  • Table 1 is a capability-level comparison of representative efficient-attention methods, using symbols for explicit, absent, and approximate capabilities.The caption defines ✓ as explicit, × as absent, and ≈ as approximate.Source: Table 1 · p. 5

4 Method

RoLA keeps a genuine cross-token global branch while making its low-rank summary compatible with relative rotary geometry. It combines outside-activation RoPE with a truncated pre-trained rotary schedule, then fuses the global correction with sparse attention.

4.1 Rotary-Positioned Low-Rank Linear Attention

  • The module pairs block-sparse attention for high-energy semantic spikes with a low-rank branch that reconstructs smooth global context.The global context is computed once per head, giving the branch O(Lrdh) cost; the design targets reusable relative-position awareness, not a lower linear-branch constant.Source: §4 · p. 5, §4.1 · p. 6

4 Method

  • RoLA’s contribution is a genuine linear-attention global branch that preserves relative geometry without an auxiliary positional module.The sparse-plus-low-rank decomposition itself is shared with prior hybrid designs.Source: §4 · p. 5

4.2 Dimensionality Truncation of the Rotary Schedule

  • The low-rank branch reuses the first r rotary coordinates in the backbone’s native ordering, with r = 64 by default and no new positional parameters.This retains r/2 complete rotary pairs and reuses the backbone’s cosine and sine tensors.Source: §4.2 · p. 6

4.3 RoPE Outside the Activation

  • RoLA applies truncated RoPE after the low-rank nonlinearity, so position-free content features can feed a reusable summary with relative-position effects at readout.Applying rotation before the activation can entangle absolute position with key features; the cited obstruction is a design motivation, not a learnability guarantee.Source: §4.3 · p. 7

4.4 Feature Alignment and Gated Fusion

  • A token-wise scalar gate combines the sparse and global outputs, with sparse-output RMS providing the scale reference for fusion.A negative gate bias starts the module close to the sparse baseline, allowing the low-rank contribution to be introduced gradually.Source: §4.4 · p. 7

4.5 Implementation Details

  • RoLA is integrated by single-stage full fine-tuning, and fused Triton kernels implement the low-rank branch using dense 2D matrix multiplications.The projections are initialized with Kaiming initialization, while the backbone is updated with a small learning rate.Source: §4.5 · p. 8

5 Experiments

Mechanism experiments test RoLA’s positional behavior and rotary-rank truncation, while generation experiments assess quality and efficiency on full-scale video models. The reported results show competitive VBench quality relative to dense attention and closer scores than pure sparse attention at matched sparsity.

5.1 Mechanism-Level Validation

  • Mechanism validation uses pre-trained Wan2.1-T2V-1.3B features from 480p, 81-frame denoising trajectories, with r = 64 and pre-RMSNorm logits.Features are normalized and averaged over heads and denoising steps.Source: §5.1 · p. 8
  • The outside-activation low-rank map closely matches the full-attention banded distance-decay pattern; inside activation distorts it, while removing RoPE yields a nearly flat map.The comparison uses a contiguous temporal slice and content-controlled inputs.Source: §5.1 · p. 8
  • Only outside-activation RoPE combines near-zero shift residual (1.8×10−5) with offset range 0.47; inside-activation RoPE has shift residual 3.7×10−1.No RoPE has zero shift residual but also zero offset range, so it is position-blind.Source: Table 2 · p. 8, §5.1 · p. 9
  • At r = 64, centered cosine similarity reaches 0.93 on both temporal and height axes, and normalized L2 error is near zero.Operator fidelity saturates at this rank; downstream quality shows clear gains through r = 64 and only marginal change beyond it.Source: §5.1 · p. 9, §5.3 · p. 11

5.2 Generation Quality

  • Generation evaluation covers Wan2.1-1.3B and Wan2.1-14B with VBench dimensions, attention sparsity, DiT denoising latency, and complementary frame-level metrics.Wan2.1-1.3B uses 480p inputs and Wan2.1-14B uses 720p inputs with 81 frames; latency excludes text encoding and VAE stages.Source: §5.2 · pp. 9–10, Table 4 · p. 10
  • RoLA reduces computational overhead while remaining competitive with dense attention on reported VBench dimensions, and its matched-sparsity scores are closer to dense attention than VSA’s.The authors report a small VBench degradation relative to full attention and a favorable quality–efficiency trade-off within evaluated settings.Source: §5.2 · p. 10

5.3 Ablation Study

  • Ablations use Wan2.1-1.3B at matched 90% sparsity; controlled global-branch comparisons fix sparse engine, pattern, and target sparsity, while a separate sweep evaluates RoLA across multiple sparsity levels.The outside ordering preserves relative-position structure; the inside ordering remains position-aware but leaks absolute position.Source: §5.3 · p. 11
  • Among tested activations, SiLU achieves the lowest alignment error; identity and Tanh underperform, while a smooth nonlinearity benefits reconstruction.The sparse-only baseline has rel-L2 = 0.612; projections are retrained for each activation.Source: §5.3 · p. 11, Table 6 · p. 11

6 Conclusion

RoLA addresses the compatibility of reusable linear aggregation with 3D relative rotary geometry by applying rank-truncated RoPE outside the nonlinear activation. The conclusion reports a favorable quality–efficiency trade-off under high sparsity, while noting that evaluation focuses on two Wan models.

6 Conclusion

  • RoLA applies rank-truncated 3D RoPE outside the nonlinear activation, preserving a reusable linear-time summary and relative positional structure without auxiliary absolute-position parameters.Mechanistic analysis and empirical evaluations indicate a favorable quality–efficiency trade-off under high sparsity.Source: §6 · p. 11
  • The evaluation focuses on two Wan models, and the appendix checks are design rationales rather than end-to-end generation guarantees.The appendix checks are not presented as theoretical contributions.Source: §6 · p. 11

A Full Forward Pass

  • Algorithm 1 combines the sparse branch, rotary-positioned low-rank branch, and distribution-aware gated fusion in a single-block forward pass.All new parameters are jointly optimized with the backbone.Source: A Full Forward Pass · p. 14

B Formal Sanity Checks and Design Rationale

The appendix presents structural checks and design intuitions for RoLA, not theoretical contributions or end-to-end generation guarantees. It verifies that outside-activation RoPE supports relative-position structure and reusable linear aggregation, while an inside-activation SiLU branch cannot yield an exact relative factorization.

B Formal Sanity Checks and Design Rationale

  • The appendix frames its formal results as structural checks and idealized design intuitions, not theoretical contributions or end-to-end generation guarantees.The truncation intuition assumes nested full-dimensional features, whereas the implemented low-rank projections are learned directly in the low-dimensional space.Source: B Formal Sanity Checks and Design… · p. 14
  • The retained RoPE coordinates are the first r coordinates in the backbone’s native ordering, comprising complete rotary pairs split across the three spatiotemporal axes.The notation does not assume a frequency ordering.Source: B Formal Sanity Checks and Design… · p. 14

B.1 Relative-Position Structure of the Outside-Activation Branch

  • With position-free low-rank feature maps, applying RoPE outside the activation makes pairwise interactions depend on relative offsets while preserving reusable linear aggregation.The argument uses RoPE’s orthogonality and group property as a structural sanity check.Source: §B.1 · p. 15
  • The outside-activation branch computes its global context as one query-independent sum at cost O(Lrdh), preserving linear cost.The rotations act outside the position-free feature maps and do not obstruct associativity.Source: §B.1 · p. 15

B.3 Obstruction for Inside-Activation Rotary Factorization

  • The inside-activation SiLU obstruction rules out exact relative factorization in the stated continuous-position setting, not approximate learnability.The appendix presents this as a design motivation rather than a claim about whether a trained branch can approximate relative behavior.Source: §B.3 · pp. 17, 19

B.2 Truncation-Error Intuition for the Rotary Low-Rank Branch

The truncation analysis considers an idealized nested setting in which a rank-r branch retains the first r rotary coordinates. Under a concentration assumption, it bounds the difference between full and truncated pairwise scores; this is design intuition, not a strict guarantee for the trained branch.

B.2 Truncation-Error Intuition for the Rotary Low-Rank Branch

  • The idealized argument assumes the first r coordinates contain complete rotary blocks and that the retained and discarded rotary subspaces remain invariant under rotation.The truncation rank is even, r = 2s, and the retained coordinates correspond to the first s complete 2 × 2 blocks.Source: §B.2 · pp. 15–16
  • The truncation bound relies on residual feature components outside the retained rotary subspace being small, as assumed for the global background branch.The assumption specifies bounds for the residual components of both feature maps.Source: §B.2 · p. 16
  • The retained subspace choice is supported by an empirical fidelity sweep rather than a frequency-ordering assumption, with Fig. 4 validating sufficiency at r = 64.The rotary schedule is geometrically spaced, and the implementation retains the first r coordinates in the backbone’s native ordering.Source: §B.2 · p. 16
  • In the idealized nested setting, retaining the first r rotary coordinates yields a bound on full-versus-truncated pairwise scores for all token pairs.The actual implementation learns Pq and Pk directly in the low-dimensional space, so this analysis is not a strict guarantee for the trained branch.Source: §B.2 · p. 16
  • Under Assumption B.2, the score error is bounded by the product of the two feature-map residual bounds.The passage gives this bound as ϵΦ(r)ϵΨ(r).Source: §B.2 · pp. 16–17

C Additional Wan2.2 and VBench 2.0 Results

This section extends evaluation to Wan2.2-T2V and selected VBench 2.0 dimensions on Wan2.1-14B and Wan2.2-A14B-T2V. The supplied passages specify evaluation settings and metrics but not comparative outcomes.

C Additional Wan2.2 and VBench 2.0 Results

D Additional Ablation Results

The additional ablations control the sparse engine and routing while varying the global branch or sparsity. The accompanying optimization discussion motivates RoLA’s single-stage fine-tuning while limiting that choice to this architecture.

E.1 Why Single-Stage Fine-Tuning Replaces the Prior Two-Stage Recipe

  • RoLA’s added parameters sit on the final low-rank attention path and fuse with the sparse branch in the inference computation, unlike the coordinate-MLP surrogate and absolute positional module used by RoPeSLR [14].The passage says RoPeSLR used a two-stage post-training procedure for its low-rank Fourier compensator.Source: §E.1 · p. 20
  • A nearly closed initial gate limits early perturbation, while the projections, gate, and backbone attention statistics must co-adapt under the final diffusion objective.The authors argue separate training would create a stage mismatch; RoLA also avoids a separate positional module whose stabilization motivates warm-up in the RoPeSLR [14] variant.Source: §E.1 · pp. 20–21
  • The dense-attention control on Wan2.1-1.3B at 480p compares pretrained full attention with the same single-stage fine-tuning protocol used for RoLA.The control keeps dense attention and matches adaptation data, batch size, and optimization steps to RoLA.Source: Table 9 · p. 21

D Additional Ablation Results

  • In a controlled global-branch ablation on Wan2.1-1.3B at 90% sparsity, the sparse engine and routing pattern are fixed across rows, and only the global branch changes.Separately, RoLA is evaluated at multiple sparsity levels.Source: D Additional Ablation Results · p. 20, Table 10 · p. 21

E.1 Why Single-Stage Fine-Tuning Replaces the Prior Two-Stage Recipe

  • The authors frame single-stage full fine-tuning as architecture-matched, not as evidence that two-stage optimization is inferior for all sparse–low-rank designs.Table 12 reports a controlled schedule comparison under the same VBench setting.Source: §E.1 · p. 21

E.2 Training-Strategy Ablation

The training-strategy ablation compares single-stage and two-stage post-training under matched experimental conditions. RoLA’s fine-tuning uses higher learning rates for newly introduced parameters and a fixed inherited rotary schedule.

E.2 Training-Strategy Ablation

  • Table 12 compares RoLA’s joint single-stage optimization with a RoPeSLR-style warm-up-then-joint schedule while holding backbone, sparsity, data, optimizer, and total update budget fixed.Both strategies are evaluated under the same VBench setting.Source: §E.2 · p. 21
  • The new parameters Pq, Pk, Wg, and bg receive a higher learning rate than the lightly adapted backbone.They are trained from initialization, while the backbone is only lightly adapted.Source: §E.2 · p. 21
  • RoLA inherits and fixes the pretrained backbone’s rotary pairing and frequency schedule, and 2,000 training samples over four epochs suffice for effective adaptation in this setting.The stage adapts a pretrained video DiT rather than learning video generation from scratch.Source: §E.2 · p. 21

F Computational Complexity Analysis

Under fixed sparsity and low-rank width, RoLA’s low-rank branch adds asymptotically negligible overhead to the sparse module. Qualitative comparisons report better temporal coherence and scene layout under high sparsity.

F Computational Complexity Analysis

  • RoLA forms global context once at O(Lrdh), reuses it across queries, and adds no positional parameters by reusing backbone rotary tensors.The rank-truncated rotation adds O(Lr) elementwise operations.Source: F Computational Complexity Analysis · pp. 21–22
  • The low-rank branch remains linear in sequence length, costs less than a full-width linear branch when r < dh/2, and matches its leading matmul FLOPs at r = dh/2.Its cost is controlled by the bottleneck width.Source: F Computational Complexity Analysis · p. 22
  • At the default r = 64 and dh = 128, RoLA and a full-width linear branch have equal leading-order matmul FLOPs.The stated advantage is preserving relative-position structure within a compressed rotary feature space, not reducing this leading-order cost.Source: F Computational Complexity Analysis · p. 23
  • Under fixed sparsity and low-rank width, the low-rank branch’s overhead is asymptotically negligible relative to sparse attention as sequence length grows.The sparse branch remains the dominant cost in this regime.Source: F Computational Complexity Analysis · pp. 22–23

G Additional Qualitative Comparisons

  • Qualitative comparisons report better temporal coherence and scene layout for RoLA under high sparsity than for competing sparse-global methods.Figure 6 shows comparisons across five video prompts.Source: G Additional Qualitative Comparisons · pp. 23–24
Loading 2609.06712v3…