Test-Time Training Done Right

We greatly appreciate the reviewer's thoughtful feedback. Below, we clarify and address each point raised:

Reviewer comment:

Clarity of presentation can be improved as the paper currently requires a very good familiarity with TTT.

Response:

1. Move the related work section on Test-Time Training (TTT) from the appendix into the main paper, aiding readers less familiar with TTT.

2. Include clear cross-references and direct links between the main paper and appendix sections, as also recommended by reviewer JVw8.

3. Move critical ablations (e.g., Figures 7 and 8) from the appendix to the main body, while relocating extensive technical details into supplementary material to streamline readability.

Reviewer comment:

The maximum number of TTT parameters tested is about 40% of the total number of model parameters. The authors should provide experiments to show the effect of further increases in the number of parameters or discuss the potential of this change.

Response: We thank the reviewer for highlighting this promising direction. Our current experiments already substantially expand TTT parameter scales beyond existing methods:

For example, the DeltaNet baseline (760M parameters) uses state sizes about 0.7% of total parameters per block, whereas our largest experiments reach up to 40%.

Figure 7(a) clearly demonstrates performance gains with increased state sizes, where novel view synthesis improves consistently from 4.4% to 40% TTT parameters, and language modeling tasks improve from 3% to 11%.

We agree that exploring even larger state sizes remains exciting and significant for future research.

Reviewer comment:

Suggestions to increase clarity: -Fig. 3 is not clear. I agree that the decoupling between the update and apply functions allows for interesting combinations but what do the blue squares indicate and how are they linked to the architectures shown above?

Response: We appreciate the reviewer's suggestion. Yes, it's a little ambiguous.

The small blue squares in Fig. 3 represent valid attention mask positions, where rows correspond to query tokens and columns to key tokens. A blue square at [row, column] indicates the column-th key token is visible to the row-th query token. We will update the figure caption accordingly for clarity.

We will adjust the caption of the figure to explain the meaning of the blue square.

Reviewer comment:

In the video diffusion experiment, the frozen weights are used for denoising. Please clarify the role played by the TTT weights.

Response: The TTT weights specifically store temporary memory from previously denoised (clean) video frames. During inference, noisy video chunks query these frozen TTT weights to retrieve contextual information from earlier frames. After the full denoising process for each chunk of video frames, TTT weights update based on the most recent clean chunk, effectively functioning as block-wise causal memory.

Reviewer comment:

The L2 norm equation for weight update should be introduced earlier in 2.1

Response: We believe the current positioning of the L2 norm equation in Section 3.2 remains appropriate. Section 2.1 references original equations from the previous TTT literature, whereas Section 3 explicitly introduces our LaCT-specific design involving weight normalization.

Reviewer comment:

Mention the GPU resources used

Response: We mentioned the GPU resources for major experiments in each modality in Section E of the Supplementary material. For reference:

For novel view synthesis: The training is completed with 64 A100 GPUs. The pre-training takes around 8 days, and each fine-tuning stage is about 12hours (thus 2 days in total).

For language model experiments

• 760M parameters scale experiment: Each experiment ran on 32 A100-40GB SXM GPUs for approximately 20 hours.
• 3B parameters scale experiments: . Each experiment ran on 64 A100-40GB SXM GPUs for approximately 50-60 hours.

For autoregressive video generation experiments:

• 1.3 B parameter scale experiments: 64 A100 80GB SXM GPUs for around 35 hours
• 14B parameter scale experiments: the 5-seconds video training takes approximately 110 hours on 64 A100 80GB SXM, then fintuned for 8.8 secs video for an additionally 13 hours on 64 H100 80GB SXM.

We appreciate the detailed review and provide clarifications and responses below.

Reviewer Concern 1:

"The core contribution appears as minor hyperparameter changes..."

Response: We respectfully clarify that our contribution extends significantly beyond minor hyperparameter modifications. Previous TTT methods (e.g., TTT-linear, TTT-MLP) have severe practical limitations. For instance, the original open-source TTT-linear implementation requires 32 A100 GPUs for 95 days for our smallest language modeling experiment (760M parameters, 40B tokens), whereas LaCT completes this in ~20 hours on identical hardware.

The primary contribution is our proposed LaCT framework, which enables efficient experimentation across diverse design spaces—including test-time optimizers, fast weight architectures, loss functions, and modalities—without cumbersome custom kernels or extensive modifications. Unlike approaches such as GLA or DeltaNet, which heavily depend on advanced Triton kernels requiring specialized expertise(which only a small group of researchers can write high-quality triton code), LaCT's pure PyTorch implementation significantly democratizes access and accelerates research in efficient nonlinear recurrent models.

Reviewer Concern 2 & 3:

"Claims about code sharing and reproducibility seem poorly justified or false."

Response: We will publicly release our code regardless of acceptance. Due to the ban of anonymous github links in rebuttal, we can paste some key part of code in chat if necessary.

We have internally prepared an open-source implementations covering:

1. Novel view synthesis model, training and inference code, with checkpoints for both object and scene dataset.

2. Language model code with huggingface transformer style (which is compitable with lot's of open-source trainer).

3. AutoRegressive Video diffusion model, training and inference code.

4. A minimal implementation of model code rapid prototyping and experimentation. 

Reviewer Concern 4(a):

"The central claim that the proposed method improves "improves hardware utilization by orders of magnitude" in the Appendix (e.g., lines 15-16) seems to be entirely dependent on the specific hardware (an A100 GPU in this case). This claim should be properly contextualized to avoid potentially misleading a non-familiar audience."

Response: We tested throughput on both A100 and H100 GPUs (see Figure 1(a)). We will contextualize hardware specifics clearly in the revision.

Importantly, our PyTorch implementation is GPU-agnostic, requiring no device-specific customization, suggesting broad applicability across modern GPUs.

Reviewer Concern 4(b):

"Other claims in the paper, such as "Language Modeling. Our model achieves competitive performance compared to SoTA methods such as DeltaNet," (lines 62-63) also seem potentially misleading."

Response: We will adjust the wording from "competitive performance compared to SoTA methods" to "competitive performance compared to current popular methods like DeltaNet."

Reviewer Concern 5(a):

"Critical details deferred to Appendix without specific references."

Response: We appreciate this feedback and will add explicit section references and hyperlinks in the revised manuscript to clearly guide readers to relevant details.

Reviewer Concern 5(b):

"There seems to be no lines of Python code in the appendix."

Response: We already included Python-style pseudocode in Section C (pp. 15-16) of the submitted supplementary material. (Please unzip the supplementary material to see the pdf named “supplementary_materials_TTTDR_neurips.pdf”. 

Reviewer Concern 6:

"Lack of proper ablations and small-scale analyses."

**Response: **

Our paper already provides several targeted ablation studies isolating key design choices on both novel view synthesis and language modeling tasks:

• Figure 7(a): Larger state size consistent boost performance

• Figure 7(b): Advanced test-time training optimizers like Muon and Momentum boost performance.

• Figure 8(a): Nonlinear fast-weight functions outperforms linear fast weight

• Figure 8(b): Evaluation of chunk recurrence versus per-token recurrence on different data structure.

Reviewer Question:

"The original TTT implementation [1] uses JAX compilation for large batch sizes and reports comparable FLOPs with linear attention. What am I missing?"

Response: The reviewer references Table 1 from [1], where TTT-linear FLOPs are 1.1× linear attention, and TTT-MLP is 1.5× linear attention. However, FLOPs alone poorly reflect real-world GPU runtime performance. Real runtime speed heavily depends on arithmetic intensity and parallelism, not just FLOPs.

For example, in our experiments, although TTT-linear theoretically has fewer FLOPs than Transformers, its runtime is around 90x longer  (90 days on 32 A100 GPUs) versus Transformers (within 1 day, same hardware). LaCT resolves this severe practical inefficiency, dramatically reducing the test-time training runtime to ~20 hours and enabling feasible, extensive experimentation.

Moreover, while JAX compilation on TPUs can fuse kernels and reduce memory overhead, achieving orders-of-magnitude speedup is uncommon. Additionally, most researchers, especially within the open-source community, predominantly use PyTorch and GPUs, environments that LaCT directly supports.

I would like to thank the authors for their rebuttal. Unfortunately, I still find that most of the criticism I raised in my review remains valid.

"We respectfully clarify that our contribution extends significantly beyond minor hyperparameter modifications. Previous TTT methods (e.g., TTT-linear, TTT-MLP) have severe practical limitations. For instance, the original open-source TTT-linear implementation requires 32 A100 GPUs for 95 days for our smallest language modeling experiment (760M parameters, 40B tokens), whereas LaCT completes this in ~20 hours on identical hardware."

I find this statement once again potentially misleading. Please correct me if I am looking at the wrong GitHub codebase, but the Pytorch TTT implementation shared by the original authors has not been optimized, has not been used for training, and is only meant for reference purposes. In fact, at the top of the README, they even state:

"The original Pytorch implementation of TTT was never optimized for inference [...] We do not recommend training with this codebase, because it is written in pure PyTorch without any systems optimization, so training will be slow, especially when the per-device batch size is small.

For training code, or to replicate results from our paper, please view our JAX codebase. For inference kernels, or to replicate speed benchmarks from our paper, please view our kernel implementations."

Beyond the hyperparameters, if the technical innovation is the "faster" Pytorch re-implementation (which was not shared at the time of submission), I do not think this matches the paper's claimed contributions.

As I mentioned in my review, I think some applications of the TTT framework, such as novel view synthesis, are interesting and should be expanded further. However, both the original paper and the rebuttal claims appear quite significantly inconsistent with the paper's actual contribution. Thus, I cannot recommend the paper for acceptance.

@torch.compile()
def zeropower_via_newtonschulz5(G):
    """
    The code is most copied from xxx [url link removed due to rebuttal constraints].
    Args:
        G: [b, d, d']
    Returns:
        X: [b, d, d']
    """
    assert len(G.shape) == 3
    X = G.bfloat16()
    if G.size(1) > G.size(2):
        X = X.transpose(1, 2)
    X = X / (X.norm(dim=(1, 2), keepdim=True) + 1e-7)
    for a, b, c in [
        (4.0848, -6.8946, 2.9270),
        (3.9505, -6.3029, 2.6377),
        (3.7418, -5.5913, 2.3037),
        (2.8769, -3.1427, 1.2046),
        (2.8366, -3.0525, 1.2012),
    ]:
        A = X @ X.transpose(1, 2)
        B = b * A + c * A @ A
        X = a * X + B @ X

    if G.size(1) > G.size(2):
        X = X.transpose(1, 2)
    return X


@torch.compile
def block_causal_lact_swiglu(
    w0: torch.Tensor,  # [b, dh, dk]
    w1: torch.Tensor,  # [b, dv, dh]
    w2: torch.Tensor,  # [b, dh, dk]
    q: torch.Tensor,  # [b, l, dk]
    k: torch.Tensor,  # [b, l, dk]
    v: torch.Tensor,  # [b, l, dv]
    lr0: torch.Tensor,  # [b, l, 1]
    lr1: torch.Tensor,  # [b, l, 1]
    lr2: torch.Tensor,  # [b, l, 1]
    momentum: torch.Tensor = None,  # [b, s, 1] # none means no momentum
    use_muon: bool = True,
    chunk_size: int = 2048,
):
    """
    Block causal LaCT with SwiGLU fast weight function.
        Apply then Update => Shifted Block Causal LaCT
    w0, w1, w2 are the fast weights. f(x) =  w1 @ (silu(w0 @ x) * (w2 @ x))
    """
    w0_norm = w0.norm(dim=2, keepdim=True)
    w1_norm = w1.norm(dim=2, keepdim=True)
    w2_norm = w2.norm(dim=2, keepdim=True)

    if momentum is not None:
        dw1_momentum = torch.zeros_like(w1)
        dw0_momentum = torch.zeros_like(w0)
        dw2_momentum = torch.zeros_like(w2)

    q = q.transpose(1, 2)  # [b, dk, l]
    v = v.transpose(1, 2)
    output = torch.zeros_like(v)
    e_index = 0
    seq_len = k.shape[1]
    for i in range(0, seq_len - chunk_size, chunk_size):
        s_index = i
        e_index = s_index + chunk_size

        # [b, l, dk]
        ki = k[:, s_index:e_index, :]  # bf16
        # [b, dv, l]
        vi = v[:, :, s_index:e_index]  # bf16
        # [b, dh, l]
        qi = q[:, :, s_index:e_index]
        # [b, l, d/1] fp32
        lr1i = lr1[:, s_index:e_index, :]  # [b, l, d/1] fp32
        lr2i = lr2[:, s_index:e_index, :]
        lr0i = lr0[:, s_index:e_index, :]

        # [b, dh, dk] @ [b, dk, l] -> [b, dh, l]
        h = torch.bmm(w2, qi)
        gate = F.silu(torch.bmm(w0, qi), inplace=True)
        # [b, dv, dh] @ [b, dh, l] -> [b, dv, l] -> [b, l, dv]
        output[:, :, s_index:e_index] = torch.bmm(w1, gate * h)

        ########## compute the gradient and update the fast weights
        # [b, dh, dk] @ [b, dk, l] -> [b, dh, l]
        gate_before_act = torch.bmm(w0, ki.transpose(1, 2))
        hidden_before_mul = torch.bmm(w2, ki.transpose(1, 2))
        hidden = F.silu(gate_before_act, inplace=False) * hidden_before_mul

        # [b, dh, dv] @ [b, dv, l] -> [b, dh, l]
        dhidden = torch.bmm(w1.transpose(1, 2), vi)

        dhidden_before_mul = dhidden * F.silu(gate_before_act, inplace=False)

        dgate = dhidden * hidden_before_mul
        dgate_before_act = silu_backprop(dgate, gate_before_act)

        # [b, dv, l] @ [b, l, dh] -> [b, dv, dh]
        # it's better to cast the mat to bf16 before bmm.
        dw1 = torch.bmm(vi, (hidden.transpose(1, 2) * lr1i).type_as(vi))  # [b, d, d]
        # [b, dh, l] @ [b, l, dk] -> [b, dh, dk]
        dw0 = torch.bmm(dgate_before_act, (ki * lr0i).type_as(dgate_before_act))
        dw2 = torch.bmm(dhidden_before_mul, (ki * lr2i).type_as(dhidden_before_mul))

        if momentum is not None:
            m_i = momentum[:, s_index:e_index, :]

            m_i = m_i.mean(dim=1, keepdim=True)

            dw0 = dw0 + dw0_momentum * m_i
            dw1 = dw1 + dw1_momentum * m_i
            dw2 = dw2 + dw2_momentum * m_i
            dw0_momentum = dw0
            dw1_momentum = dw1
            dw2_momentum = dw2

        if use_muon:
            dw1 = zeropower_via_newtonschulz5(dw1)
            dw0 = zeropower_via_newtonschulz5(dw0)
            dw2 = zeropower_via_newtonschulz5(dw2)

        w0 = w0 + dw0
        w1 = w1 + dw1
        w2 = w2 + dw2
        w0 = w0 / (w0.norm(dim=2, keepdim=True) + 1e-5) * w0_norm
        w1 = w1 / (w1.norm(dim=2, keepdim=True) + 1e-5) * w1_norm
        w2 = w2 / (w2.norm(dim=2, keepdim=True) + 1e-5) * w2_norm

    s_index = e_index  # update_length - mini_batch_size
    e_index = seq_len
    qi = q[:, :, s_index:e_index]
    # get the final output
    # [b, dh, dk] @ [b, dk, l] -> [b, dh, l]
    h = torch.bmm(w2, qi)
    gate = F.silu(torch.bmm(w0, qi), inplace=True)
    # [b, dv, dh] @ [b, dh, l] -> [b, dv, l]
    output[:, :, s_index:e_index] = torch.bmm(w1, gate * h)
    # [b, l, dv]
    return output.transpose(1, 2)

@torch.compile
def bidirectional_lact_swiglu(
    w0: torch.Tensor,  # [b, dh, dk]
    w1: torch.Tensor,  # [b, dv, dh]
    w2: torch.Tensor,  # [b, dh, dk]
    q: torch.Tensor,  # [b, l, dk]
    k: torch.Tensor,  # [b, l, dk]
    v: torch.Tensor,  # [b, l, dv]
    lr0: torch.Tensor,  # [b, l, 1]
    lr1: torch.Tensor,  # [b, l, 1]
    lr2: torch.Tensor,  # [b, l, 1]
    use_muon: bool = True,
) -> torch.Tensor:
    """
    Bidirectional LaCT with SwiGLU fast weight function.
    w0, w1, w2 are the fast weights. f(x) =  w1 @ (silu(w0 @ x) * (w2 @ x))

    About precision:
        w0, w1, w2 are mostly likely fp32. 
        q, k, v are fp16.
        lr0, lr1, lr2 are fp32.
        The forward, backward produce bf16 gradients, updated fast weights are fp32.
        The final output are bf16.


    FLOPS:
        (assume dk=dv denoted as D, hidden dimension of swiglu-mlp is H, ignore muon)
        Forward pass with key: 4 * D * H * L * B
        Backward pass: 8 * D * H * L * B
        Forward with Query: 6 * D * H * L * B
        Total: 18 * D * H * L * B
    Outputs:
        o: [b, l, dv]
    """

    # adding detach here sometimes improves stability.
    w0_norm = w0.norm(dim=2, keepdim=True)
    w1_norm = w1.norm(dim=2, keepdim=True)
    w2_norm = w2.norm(dim=2, keepdim=True)

    q = q.transpose(1, 2)  # [b, dk, l]
    v = v.transpose(1, 2)

    ######### update the fast weight w0, w1, w2 with test-time training #########

    #### Forward pass with key
    # [b, dh, dk] @ [b, dk, l] -> [b, dh, l]
    gate_before_act = torch.bmm(w0, k.transpose(1, 2))
    hidden_before_mul = torch.bmm(w2, k.transpose(1, 2))
    hidden = F.silu(gate_before_act, inplace=False) * hidden_before_mul

    #### Backward pass to compute fast weight gradients
    # [b, dh, dv] @ [b, dv, l] -> [b, dh, l]
    dhidden = torch.bmm(w1.transpose(1, 2), v)

    dhidden_before_mul = dhidden * F.silu(gate_before_act, inplace=False)
    dgate = dhidden * hidden_before_mul
    dgate_before_act = silu_backprop(dgate, gate_before_act)

    # [b, dv, l] @ [b, l, dh] -> [b, dv, dh]
    dw1 = torch.bmm(v, (hidden.transpose(1, 2) * lr1).type_as(v))  # [b, d, d]
    # [b, dh, l] @ [b, l, dk] -> [b, dh, dk]
    dw0 = torch.bmm(dgate_before_act, (k * lr0).type_as(dgate_before_act))
    dw2 = torch.bmm(dhidden_before_mul, (k * lr2).type_as(dhidden_before_mul))
    

    if use_muon:
        w0 = zeropower_via_newtonschulz5(dw0)
        w1 = zeropower_via_newtonschulz5(dw1)
        w2 = zeropower_via_newtonschulz5(dw2)

    w1 = w1 + dw1
    w0 = w0 + dw0
    w2 = w2 + dw2

    w0 = w0 / (w0.norm(dim=2, keepdim=True) + 1e-5) * w0_norm
    w1 = w1 / (w1.norm(dim=2, keepdim=True) + 1e-5) * w1_norm
    w2 = w2 / (w2.norm(dim=2, keepdim=True) + 1e-5) * w2_norm

    ######### apply the updated fast weights to the query #########

    # [b, dh, dk] @ [b, dk, l] -> [b, dh, l]
    h = torch.bmm(w2, q)
    gate = F.silu(torch.bmm(w0, q), inplace=True)
    # [b, dv, dh] @ [b, dh, l] -> [b, dv, l] -> [b, l, dv]
    o = torch.bmm(w1, gate * h).transpose(1, 2)

    return o

"The reviewer criticized us for not being able to provide code at the time of submission."

I never did criticize not providing code by itself. This is not a requirement of NeuriPS. I pointed out the inconsistency with the checklist and the fact that the implementation was being claimed as the core contribution.

"Claims in the checklist appear to be poorly justified or even completely false. For instance, the authors answered [YES] to the questions regarding code sharing and reproducibility. Yet, they did not provide any code with the submission, with this claim being "justified" by a promise of sharing code conditioned on acceptance." (Original review)

"I would suggest avoiding claiming that speeding up TTT N times, thanks to a new Pytorch implementation to be the paper's main contribution. I believe this is not only partially out-of-scope for a NeurIPS paper's main contribution, but it would also require at least providing reviewers with the code at the time of submission." (Reply above)

Regarding your other points, I am not sure how productive this conversation is getting, as I am starting to fear that my criticism and all the other areas of improvement I highlighted in my review and prior responses are not being taken seriously, and the authors are now overly focusing on one very specific point, the code contribution of the work.

We thank the reviewer for their detailed feedback. We will address the main concerns below.

Reviewer comment:

The authors argue that LaCT's efficiency enables the use of sophisticated optimizers but do not provide evidence that other TTT methods could not also benefit from Muon, which is a critical missing ablation.

Response: We appreciate the reviewer highlighting this point. However, a direct comparison by applying sophisticated optimizers such as Muon to original test-time training methods (e.g., ttt-linear or ttt-mlp) is practically infeasible due to severe computational limitations. For instance, in our smallest language modeling experiment (760M parameters trained on 40B tokens), the original open-source ttt-linear implementation requires 32 A100 GPUs for 95 days, and ttt-mlp requires 150 days, whereas LaCT and other baselines complete training in approximately 20 hours on identical hardware. Integrating Muon into the original ttt-linear method would further increase the computational cost (FLOPS) by approximately 30x, making such comparisons impractical (see detailed FLOPS derivation in Equation (20) of the supplementary).

Therefore, original TTT methods fundamentally cannot leverage advanced optimizers at modern scales. In contrast, our large-chunk design efficiently supports sophisticated optimizers with only dozens of lines of PyTorch code. Moreover, it enables rapid adaptation not only of test-time optimizers but also of different fast weight architectures, loss functions, and diverse data modalities efficiently as shown in the paper.

Reviewer comment:

The paper fails to clearly disentangle the benefits of the core LaCT architecture from the optimizer used.

Response: We respectfully disagree with the reviewer’s assessment. The primary novelty and contribution of our paper lie not simply in the choice of optimizer (e.g., Muon) but rather in the LaCT architectural framework itself. This architecture allows rapid and efficient experimentation across diverse design choices—including optimizers, fast weight architectures, loss functions, and data modalities—without relying on complex custom kernels,  extensive code modifications, or suffering from extremelly low hardware utilizations. 

Importantly, most deep learning researchers are unable to write cumbersome custom kernels but can readily implement PyTorch code. Thus, LaCT significantly democratizes research on test-time training and accelerates progress in next-generation efficient nonlinear recurrent neural networks.

Furthermore, our paper includes extensive ablation studies clearly isolating the contributions of different design choices on both the novel view synthesis and language model experiments:

• Figure 7(a): Impact of state size.
• Figure 7(b): Effect of advanced optimizers.
• Figure 8(a): Comparison of nonlinear versus linear fast-weight functions.
• Figure 8(b): Evaluation of chunk recurrence versus per-token recurrence.

These analyses conclusively demonstrate that nonlinear recurrent fast weights, larger state sizes, and advanced optimizers each independently enhance performance—and importantly, all these components become computationally efficient on modern GPUs only within the large-chunk architecture.

We hope the above explanation clarifies our major contributions.

---

On more benchmark evaluations

Reviewer comment:

evaluation on a suite of downstream tasks (e.g., ARC, Hellaswag) for language models

Response: We thank the reviewer for the suggestion. We have added the suite of evaluations:

As the reviewer noted, these benchmarks primarily assess the model's general commonsense knowledge. However, these metrics alone do not precisely capture the long-context modeling capability of the models. In contrast, the per-position-loss and retrieval task (S-NIAH), used in our original evaluation, specifically measure the model's ability to handle and retrieve information from longer contexts.

Reviewer comment:

lacks standard video generation metrics.

Response: We follow Yin et al[1]. to include full suite of metrics in Vbench[2]. This concern is also raised by reviewer 3KQP.

Due to character limits, we put the table of VBench results in replies of reviewer 3KQP.

[1] Yin, Tianwei, et al. "From slow bidirectional to fast autoregressive video diffusion models." CVPR2025

[2] VBench: Comprehensive Benchmark Suite for Video Generative Models. CVPR 2024.

---

On Comparison with Additional Baselines

Reviewer comment:

The paper lacks direct comparisons against other relevant baselines that combine local attention with recurrence, such as InfiniAttention or MEGA.

Response: In our language modeling experiments, the main paper already compares our method against two state-of-the-art baselines combining local attention with advanced recurrence: GLA with SWA and DeltaNet with SWA.

Here, we additionally provide a direct comparison with InfiniAttention. Since there is no official open-source implementation available, we implemented InfiniAttention ourselves following Equations (7)-(10) from their paper. Our initial implementation encountered gradient explosion issues after approximately 20K iterations. To enhance stability, we added two additional normalization layers:

• An RMS normalization layer on the output of the memory module, as is done in GLA, DeltaNet and LaCT
• Weight normalization on the linear recurrent memory as is done in LaCT.

Both normalization methods significantly improved stability for InfiniAttention.

We evaluate the models using the average per-position loss over 2.5 billion validation tokens, consistent with Figures 5(a) and 5(c). Due to rebuttal limitations on adding new figures, we present the results in the table below, showing the average token loss for each 4K-token chunk across eight chunks (total sequence length of 32,768 tokens):

We also evaluate on S-NIAH:

We have two intuitions about InfiniAttention's lower performance:

• It employs chunked window attention rather than sliding window attention, limiting its ability to effectively capture continuous context.
• Its recurrence uses a chunk-level linear recurrence with a delta rule, which typically has lower expressiveness compared to the per-token linear recurrence with delta-rule used by DeltaNet in language task.

On Paper Organization and Clarity

We thank the reviewer for the excellent suggestions on improving the paper's structure. We agree and will make the following changes in the revised version:

• Related Work: The core "Related Work" section on Test-Time Training will be moved from the appendix to the main body of the paper to better position our contributions.

• Ablation Studies: The key ablation studies (Figures 7 and 8), which are vital for understanding the method's design, will be moved into the main paper.

• Experimental Details: To create space for these additions, less critical implementation and dataset details from Section 5 will be moved to the appendix.

• Minor Corrections We thank the reviewer for catching these errors. We will correct the duplicate reference and fix all typos (including "validation" and "Muon") in the final version.

I appreciate the reviewer’s engagement. Here we address your new comments:

Reviewer's comment

Could you please elaborate your finding on VBench results. I see that the "Frame Quality" and "Text Alignment" have deteriorated in performance, right?

Response:

When comparing our method (1.3B) against Transformer SWA (1.3B) and Transformer (1.3B), our approach achieves higher scores in both Temporal Quality (0.9250 vs. 0.9114/0.9232) and Frame Quality (0.6360 vs. 0.6303/0.6219). However, there is a minor decrease in Text Alignment (0.2475 vs. 0.2477/0.2478), indicating a slight trade-off.

Reviewer's comment

Report the Total Score

Response:

Thank you for this suggestion. We have computed and provided the Total Score using the official VBench scripts, which aggregate and normalize all individual scores. Here are the results:

Reviewer's comments

Also could you clarify why the overall consistency score of the Original Wan (14B) model 0.2499 is much lower than the results in the leaderboard 0.2749?

Response:

Thanks for pointing out this discrepancy. Upon investigating this inconsistency, I found a related github issue (#163, "Inquiry about Wan2.1-T2V-14B Results on vbench (2025-07-08)") of the VBench GitHub repository. (We apologize for not being able to include direct GitHub issue links during the rebuttal period.)

According to the latest response in that issue from last week, the Original Wan model on the leaderboard was evaluated using the augmented prompts with longer text description. (my guess is that the Wan T2V API would rewrites short captions into more detailed text prompts that align closely with their training caption distribution.). Also, the leaderboard results sampled the videos for five times, and we sampled each prompt one time.

We can regenerate videos using these augmented text prompts and re-running evaluations. We will update our results accordingly if we can complete this task by August 8th (the end of the author response period).

Thank you for your clarification and reporting the total score. I believe the strongest results are obtained for this video generation task. The gains for the LLM related tasks seem insignificant or slightly worse for 760M, but still informative to report.

Overall, the new results have made the evaluations more convincing, and I will raise the rating to 3. I will engage in the reviewer's discussion and may increase the score further.

We thank the reviewer for their constructive and insightful feedback. Below we address the key questions and suggestions:

VBench Results

More perceptual evaluations for video generations

Response: We appreciate the reviewer's suggestion, which also aligns with Reviewer cwFf’s comment. To provide more comprehensive perceptual evaluations, we conducted additional experiments using the widely recognized VBench benchmark suite [2], following evaluation practices established by Yin et al. [1].

Specifically, we used 942 standard text prompts from VBench and generate videos for all compared models. For our 1.3B-parameter model, we generated two videos per prompt. For our 14B-parameter model, we generated one video per prompt due to time constraints.

Below are the summarized evaluation results using standard VBench metrics, which we will include in the revised manuscript:

---

Our method consistently performs on par or better than other autoregressive video generation baselines, regarding temporal and frame quality.

Below is the detailed metric comparison table (we would include a radar chart version in the revised version):

[1] Yin, T., et al. "From slow bidirectional to fast autoregressive video diffusion models." CVPR 2025.

[2] VBench: Comprehensive Benchmark Suite for Video Generative Models, CVPR 2024.

Understanding Muon

Why the model learning has a big performance gap when using momentum or moun based optimizer for LM (figure 5)

Response: This is an insightful question, and we have some initial intuitions.

First, some background: Muon is a novel optimizer recently adopted by leading AI labs as a replacement for AdamW in large language model pretraining (see Kimi K2 [3]). In our paper, we demonstrate the effectiveness of Muon not during pretraining but specifically during test-time training, as illustrated in Figures 5 and 7(b) on page 15.

Understanding why Muon performs effectively is an active area of ongoing research. One recent paper that significantly influenced our thinking is reference [4], which shows that spectral normalization methods on gradients like Muon can notably enhance model learning under extreme label imbalance conditions.

This intuition aligns closely with our S-NIAH results, where Muon considerably improves performance (Figure 5(b, d)). Specifically, when evaluating on S-NIAH, the model must memorize small segments of passkeys embedded within irrelevant contexts, analogous to label imbalance scenarios. Because Muon orthogonalizes and normalize the spectral of gradients, it amplifies gradient magnitudes associated with less frequent tokens, thereby facilitating the model's ability to retain long-tail information.

These insights are preliminary hypotheses and may require further validation. We intend to add this point in an appendix paragraph for revised version. We hope this perspective inspires both reviewers and future readers.

[3] Kimi K2: Open Agentic Intelligence. arxiv 2025.

[4] On Generalization of Spectral Gradient Descent: A Case Study on Imbalanced Data