Thank you for your meaningful comments! We sincerely appreciate your thoughtful feedback and hope that our responses below will help address your questions:

1. Clarification regarding tensor operations

Thank you for raising this important point! We will include a more detailed explanation in the revised version. Let us walk through a simple example to clarify here: consider an input $\mathbf{x} \in \mathbb{R}^{B \times d}$, where $B$ is the batch size. We reshape $\mathbf{x}$ into $\mathbb{R}^{B \times d_0 \times d_1}$ to match the first two TT-cores: $\mathbf{C}^0 \in \mathbb{R}^{1 \times d_0 \times r_0}$, $\mathbf{C}^1 \in \mathbb{R}^{r_0 \times d_1 \times r_1}$, within a 4-core Tensor Train: $\{\mathbf{C}^0, \mathbf{C}^1, \mathbf{C}^2, \mathbf{C}^3\}$ setting.

• The TT output $\mathbf{y}_{\text{TT}} = \mathbf{xC^0C^1C^2C^3}$ consists of two stages: (1) down-projection via $\mathbf{C}^0, \mathbf{C}^1$ and (2) up-projection via $\mathbf{C}^2, \mathbf{C}^3$. Contraction order within each stage is flexible, but for simplicity, we follow the natural order (0→1→2→3),i.e., we first contract $\mathbf{x}$ with $\mathbf{C}^0$, eliminating the $d_0$ mode and producing $\mathbf{xC^0} \in \mathbb{R}^{B \times d_1 \times r_0}$, then we contract $\mathbf{C}^1$ with $\mathbf{xC^0}$, eliminating $d_1$ and $r_0$ and producing $\mathbf{xC^0C^1} \in \mathbb{R}^{B \times r_1}$.
• Each core $\mathbf{C}^i$ is contracted with intermediate output $\mathbf{xC^0\cdots C^{i-1}}$, so there is no case where multiple cores are contracted with the same $\mathbf{x}$ simultaneously.
• To illustrate intra-layer LaX, we continue from $\mathbf{xC^0C^1} \in \mathbb{R}^{B \times r_1}$ and contract it with $\mathbf{C}^2$, yielding $\mathbf{xC^0C^1C^2} \in \mathbb{R}^{B \times d_2 \times r_2}$. If the shapes align (i.e., $d_1 r_0 = d_2 r_2$), we can fuse $\mathbf{xC^0}\in \mathbb{R}^{B \times d_1 \times r_0}$ with $\mathbf{xC^0C^1C^2}$ directly to build a residual path. If not, we use a gating mechanism to align and merge the features.

2. Additional experiments and discussions on LaX with tensor decomposition

Thank you for your insightful feedback. We have conducted additional experiments specifically targeting Tensor Train models, which we summarize below and plan to include in the revised version.

2.1 LaX with CoLA Activations in Tensor Train Models

We explored how LaX interacts with CoLA-style activations in TT models by injecting a nonlinearity after the down projection result $\mathbf{xC^0C^1}$. This extension is conceptually natural. Notably, combining LaX with CoLA leads to a synergistic effect, outperforming both standard TT models and SVD-based low-rank baselines:

This result may suggest that LaX can amplify the benefits of CoLA-style activations in tensor format models, significantly boosting the weakest base model to be even stronger than the full-rank model.

2.2 When Intra-LaX Helps: Pros and Cons

Understanding where Intra-LaX contributes most is crucial, especially for TT-LaX models with more cores. To investigate this, we conducted a series of ablation experiments, where ✓ indicates that LaX is applied to all layer types:

These results reveal two key findings:

1. Inter-LaX plays a relatively primary role in performance gains for TT models.
2. When applying Intra-LaX selectively, MLP blocks benefit more than Attention blocks.

A more granular study, e.g., dissecting Q, K, V projections, is still needed to pinpoint where LaX is most critical. We plan to further investigate this in future work.

We sincerely thank the reviewer for their thoughtful and constructive feedback. We would like to take this opportunity to clarify that LaX is not designed solely for Adapters; rather, it is a general mechanism for enhancing low-rank models, applicable to both low-rank foundation model pre-training from scratch and adapter-based fine-tuning scenarios. We hope that our point-by-point responses below will address any remaining concerns and further clarify the scope and impact of our contributions.

1. Clarification of LaX's broader impact and theoretical justification

Thank you for the helpful comments. We provide clarifications on the two points raised:

• We address a major challenge in low-rank pre-training. Pre-training low-rank foundation models is fundamentally different from LoRA-style fine-tuning. Rather than inserting adapters, these models replace dense layers entirely with low-rank layers and are trained from scratch. We would like to emphasize that enabling effective low-rank model pre-training is a key contribution of LaX. While such low-rank models typically suffer from performance degradation due to limited capacity, LaX consistently helps recover, and even surpass, the performance of full-rank models (Table 1, 2, 5, 6), using 2–3× fewer parameters overall and with negligible computational overhead.
• Theoretical analysis: As stated in the paper, our primary motivation is to improve the capacity of low-rank layers without explicitly increasing their actual matrix rank. However, we do offer a theoretical justification, complemented by practical evidence, showing that LaX increases the effective rank[1][2][3][4] after being applied to the SVD layer. Specifically:

This expression is derived from concentration results of the Frobenius and spectral norms. In the simulation, we train a simple SVD/LaX layers with identical encoders and decoders on the CIFAR-10 dataset. Due to rebuttal constraints, we are only able to present how the amplification ratio $\frac{\mathrm{sr}(\mathbf{W}^{\text{LaX}})}{\mathrm{sr}(\mathbf{W}^{\text{SVD}})}=\left( \frac{\sqrt{2d_{\text{out}}} + \sqrt{2d_{\text{in}}}}{\sqrt{d_{\text{out}}} + \sqrt{2d_{\text{in}}}} \right)^2$evolves during training in tabular format. We can see that SVD w/ LaX always has a higher effective rank than w/o LaX.

2. Comparison with other strong baseline on LLM fine-tuning

Thank you for pointing this out. We agree that integrating LaX with more advanced methods is a meaningful direction.

• We would like to clarify that we intentionally chose standard mainsteam methods in both pre-training and fine-tuning in order to isolate and clearly demonstrate the effect of LaX without interference from additional techniques.
• Regarding other PEFT baselines/adaptive rank methods, LaX is technically compatible with many existing approaches. For example, in AdaLoRA, the LaX Gate can align low-rank spaces with varying shapes. Notably, while existing adaptive low-rank pre-training methods typically fall short of full-rank performance, LaX enables vanilla low-rank models to surpass the full-rank baseline. Due to the limited time window, we focused on one strong PEFT baseline, DoRA [5], and show the comparative results below. DoRA + LaX consistently performs better or same than plain DoRA in 7 out of 8 tasks.

The experiment is conducted on Llama-7B with an adapter rank=16, using the same hyperparameters as in [5].

References:

[1] Vershynin R. High-dimensional probability: An introduction with applications in data science[M]. Cambridge university press, 2018.

[2] Tropp J A. An introduction to matrix concentration inequalities[J]. Foundations and Trends® in Machine Learning, 2015, 8(1-2): 1-230.

[3] Rudelson M, Vershynin R. Sampling from large matrices: An approach through geometric functional analysis[J]. Journal of the ACM (JACM), 2007, 54(4): 21-es.

[4] Ipsen I C F, Saibaba A K. Stable rank and intrinsic dimension of real and complex matrices[J]. arXiv preprint arXiv:2407.21594, 2024.

[5] Liu S Y, Wang C Y, Yin H, et al. Dora: Weight-decomposed low-rank adaptation[C]//Forty-first International Conference on Machine Learning. 2024.

Dear Reviewer 1Hwt,

Thanks a lot for providing valuable comments to our paper.

We understand that you may have a very busy schedule at this moment. As the extended deadline of the author-reviewer discussion is approaching, we would like to send you a gentle reminder according to the recent NeuRIPS email instruction.

We would highly appreciate it if you can look at our rebuttal and let us know whether your concerns have been addressed or not.

Sincerely,

The authors

We are very grateful with the reviewer’s insightful review. Hope our point-wise responses would further clear things out:

1. Confidence bounds and more training curves.

We agree that including confidence bounds would provide additional insights. Due to the long pre-training time and limited rebuttal period, running multiple trials for all models (which also include related methods such as ReLoRA, GaLore, SLTrain, LORO, etc) to estimate variance is infeasible. However, to provide a valid proof of concept, we have ran 60M LLaMA CoLA+LaX with 5 random seed, and estimate the 95% confidence interval as [33.15, 33.36].

2. Where LaX benefits mostly.

Thank you for the insightful question. To assess where LaX is most effective, we analyzed the checkpoint from Table 2 (rank=256), focusing on the scalar values of the Linear Gates—with higher values indicating greater contribution from LaX.

Our analysis reveals several interesting observations:

• For the Q and K projections, LaX contributes more in the first half of the layers than the latter half, with a sharp increase in gate values observed in the final layer.
• The V projection exhibits a distinct pattern, where LaX contributes less in the earlier layers compared to the later ones.
• In the attention output projection layer, only Layers 1, 10, and 11 show relatively higher gate values, while the remaining layers contribute minimally.
• For FC1 in the MLP, the contribution of LaX is more evenly distributed across layers, but still follows the general trend of higher values in the early and final layers.
• For FC2 in the MLP, only Layers 4 and 11 show meaningful gate values, with others being close to zero.

We thank the reviewer once again for their insightful observations, which may point toward a promising future direction—a sparse variant of LaX that selectively activates only the most impactful layers.

3. How does LaX interact with other techniques for improving low-rank training?

We fully agree that exploring the interaction between LaX and other techniques for low-rank training is a highly promising direction. In line with your suggestion regarding Q-LoRA, we’d like to share some preliminary results from an ongoing work investigating the integration of LaX with quantization methods.

• We pre-train ViT-B and its CoLA + LaX variant (w/ Tensor Gate) at FP32 on Imagenet-1k, then we perform a series of simple post-training quantization (Optimum Quanto from Huggingface for linear quantization) on the resulting model,

• As shown, LaX integrates effectively with quantization techniques. Notably, Int4 CoLA + LaX achieves performance comparable to the full-precision dense ViT-B model, while achieving a 16× reduction in model size (2× in parameters and 8× in precision).

4. Ablation for applying LaX to different layer types and for Normalization.

We conducted the suggested ablation experiments specifically on the pre-training of the SVD-ViT-B model on ImageNet-1K.

• Ablation for applying LaX to different layer types:

As shown, applying LaX to the V-proj yields the most notable gains among attention layers. Between MLP and Attention Blocks, the latter benefits slightly more. Notably, the benefit of LaX appears additive, with multiple applications leading to cumulative improvements.

• Ablation for Normalization:

Removing LayerNorm (LN) from LaX causes training to fail from the start, as expected. LN stabilizes activation distributions and prevents gradient instability; without it, training collapses due to shifting statistics.

5. Dependency on baseline quality. Modest improvements in some cases. The cost-benefit trade-off isn't always clear.

Thank you for your valuable comments. As discussed in the paper:

• LaX is specifically designed to recover the performance lost in low-rank models relative to their full-rank counterparts. Naturally, when starting from a strong low-rank baseline, such as CoLA, the improvement is less significant. Nonetheless, the integration of LaX consistently enables CoLA to surpass the full-rank baseline in both vision and language modeling tasks.
• LaX introduces minimal computational overhead, even with the most expensive practical variant. As noted in the practical guideline section, Tensor Gate is the most costly one in practice. The relative FLOPs overhead, compared to the baseline low-rank model,

6. The optimal gate type appears to be task and model dependent, requiring empirical search rather than principled selection.

We appreciate the reviewer’s concern and agree that different gate types yield different trade-offs. However, we would like to clarify that the core motivation of the gate design in LaX is to align latent dimensions across or within low-rank layers, which is not mandatory in many cases, especially with SVD and CoLA. Without gating, LaX still yields consistent gains across various domains, as shown in Table 2 and Table 6.

7. Novelty/theoretical analysis.

Thank you for the thoughtful comment. While LaX incorporates residual connections, its core novelty lies in addressing a unique challenge in low-rank adaptation. As stated in the paper, our primary motivation is to enhance the capacity of low-rank layers without increasing their matrix rank.

• Importantly, the effective rank $\mathrm{sr}=\frac{\mathbf{||W||}^2_F}{||\mathbf{W}||^2_2}$ of a layer is often lower than its matrixl rank. This observation opens up an opportunity: improve the effective rank in a lightweight and principled way — latent feature fusion.
• We support this intuition with both theoretical justification under assumptions and empirical evidence during the actual training process, demonstrating that LaX increases the effective rank and leads to consistent performance gains. While the mechanism may resemble residual connections in form, LaX is functionally and conceptually distinct in its purpose. Specifically (notion is the same as in our paper),

• In the simulation, we train a simple SVD/LaX layers with identical encoders and decoders on the CIFAR-10. Due to rebuttal constraints, we are only able to present how the amplification ratio $\frac{\mathrm{sr}(\mathbf{W}^{\text{LaX}})}{\mathrm{sr}(\mathbf{W}^{\text{SVD}})}$and training loss evolve during training in tabular format. We can see that SVD w/ LaX always has a higher effective rank than w/o LaX.

8. Bigger models like ViT-H/G.

We agree that studying the scaling behavior of LaX combined with low-rank models is an important direction. Due to budget constraints, we were unable to pre-train many large models. But we have included four pre-training results on 1B-scale LLMs in Table 5, which serve as meaningful references for a scale positioned between ViT-H (~650M) and ViT-G (1.8B). At 1B scale, LaX still consistently brings benefits to all the low-rank models.

9. More LoRA-like variants.

For low-rank fine-tuning, we adopt LoRA as the base method because it is the foundational approach in this domain and we aim to isolate the impact of LaX itself, just as in our pre-training experiments. However, in order to explore LaX’s orthogonal benefits when combining with other modern LoRA variants, we added addtional experiments on DoRA [1], and show results below. Notably, DoRA + LaX consistently performs better or same than plain DoRA in 7 out of 8 tasks.

The experiment is conducted on Llama-7B with an adapter rank=16, using the same hyperparameters as in [1].

References: [1] Dora: Weight-decomposed low-rank adaptation.

We really appreciate the reviewers’ thoughtful review and recognition of our contributions. Hope our point-wise responses would further clear things out:

1. Clarification of why CoLA+LaX outperforming full-rank baselines

Thanks for your comments. We confirm that CoLA + LaX outperforms the dense model in both Vision (Table 1, we train both full-rank and low-rank models using configurations from [1]) and Language model pre-training (Table 5, we directly compare LaX with baselines reported in related work [Galore, SLTrain, CoLA, LORO]). According to the CoLA paper, CoLA improves architectural efficiency by replacing an arbitrary linear operator by a composition of linear-nonlinear-linear operator, which is theoretically capable of approximating any continuous function, powered by the universal approximation theorem. We speculate that such a deeper, nested structure might empower CoLA with a higher capacity while using fewer parameters, and LaX further boosts its performance, therefore surpassing full-rank models.

2. Clarification of the parameters introduced by LaX’s layer norm

Thanks! We would like to clarify that the LayerNorm modules introduced in LaX indeed introduce parameters to the base model. However, the increase is negligible — only $2 r$ parameters per layer, $\frac{1}{d}$ of the base model. Given that $d$ is normally very large, the overall increase is trivial and does not visibly change the reported parameter counts due to rounding. We appreciate the opportunity to clarify this point and will consider explicitly noting it to avoid confusion.

3. Do authors think CoLA+LaX can still outperform full-rank models if using other optimizers, such as Muon?

Thank you for your insightful comment! It is possible that a fine-tuned optimizer might further improve the accuracy of either CoLA+LaX, full-rank models or the both. In this paper, we demonstrate the performance boost of LaX while using mainsteam optimizer such as Adam. Definitely many other research opportunities can be inspired by considering the combinations of various optimizers and architectures.

4. Discussion regarding ResLoRA/LoR2C.

We thank the reviewer for highlighting these relevant and recent works. Both are targeting the gradient vanishing/explosion challenges in PEFT. We acknowledge their potential relevance and see promising opportunities for incorporating their insights into future extensions of our approach.

• ResLoRA tackles gradient propagation issues in LoRA by introducing shortcut connections at various positions within the adapter structure, supported by a mathematical analysis from an optimization perspective. LoR2C mitigates gradient vanishing by introducing residual connections across entire blocks with an extra LoRA block.
• In LaX, we currently apply LaX between the same type of layers across transformer blocks, guided by a prior on semantic alignment. However, their work suggests that exploring LaX variants that span different connection types is possible, particularly in the context of Tensor Train models. Owing to their richer and more structured latent spaces, Tensor Train models offer greater flexibility for experimenting with alternative LaX connection patterns. We are looking forward to exploring this in our future research.
• Moreover, we are also deeply interested in the theoretical understanding of LaX. We believe that their optimization-based analyses could inspire new theoretical analyses to explain LaX.

We thank the reviewer once again for pointing out these valuable works, which may serve as complementary directions to LaX in our future research efforts!

References:

[1] Dosovitskiy A, Beyer L, Kolesnikov A, et al. An image is worth 16x16 words: Transformers for image recognition at scale[J]. arXiv preprint arXiv:2010.11929, 2020.

Thank you for your response! And we are glad to hear that we have resolved almost all your concerns. We further clarfiy our hyper-parameter configurations as below:

• ViT Pre-Training Experiments: As we mentioned in our rebuttal response, we adopted all the hyper-parameters from [1] except the learning rate. The original LR=3e-3 caused unstable training in LaX, therefore we lowered it to 1e-3 and kept this setting for all the experiments.
• LLM Pre-Training Experiments: As we mentioned in our rebuttal response, all baselines results are directly reported from related works (GaLore, SLTrain, CoLA, LORO). And we adopted most hyper-parameters from these baseline methods, except learning rate and gradient clipping scale.
  • For gradient clipping, we simply found using a slightly smaller scale, i.e., 0.5 instead of 1.0, yields relatively better training stabilities, therefore we used this setting for all LLM pre-training experiments.
  • For learning rate, we only swept for small models, such as 60M and 130M. And our early explorations at 60M scale suggested LaX could benefit from significantly larger learning rate, typically >1e-2. Then we tried gradually increasing it until the model diverges, which gave us 4e-2 for 60M scale. Then larger scale models will start from this LR, if it diverges, we typically just halve the value, and it would yield a very good setup. If not, then we further halve the value. This is how we found the LR for 130M, 350M and 1B. Sweeping would be too expensive for models >=350M.
• LLM Fine-Tuning Experiments: We simply followed the same setup in [2], without any hyper-parameter tuning or sweeping.

References:

[1] Dosovitskiy A, Beyer L, Kolesnikov A, et al. An image is worth 16x16 words: Transformers for image recognition at scale[J]. arXiv preprint arXiv:2010.11929, 2020.

[2] Hu, Zhiqiang, et al. "Llm-adapters: An adapter family for parameter-efficient fine-tuning of large language models." arXiv preprint arXiv:2304.01933 (2023).

Thank you for the reply, you have addressed almost all my concerns.

As stated in my review, Its quite important for reproducibility to report the values of hyperparameter configurations swept for each method.

Could the authors provide this information?