On the Duality between Gradient Transformations and Adapters

Thank you for your review! We address your questions below:

Key points

Is it possible to prove a similar equivalence where you use the moore-penrose pseudo-inverse to project back up?

Great question! It is possible that an equivalence exists if one is not using a LoRA adapter. For example, one could imagine defining a PyTorch primitive function whose forward pass behaves like a regular (one-sided) LoRA adapter $S^\top A$, but whose backward pass computes the gradient of $A$ using the Moore-Penrose inverse instead of $S^\top$. In principle this should be equivalent to training with GaLore (with a Moore-Penrose up-projection) in the optimization trajectory sense, though we are not quite sure what the mathematical interpretation of this would be.

The comparison between various adapter and gradient transformation methods in Table 2 was interesting, though I wasn't sure how this was directly relevant to the equivalence they proved.

The main aim of Table 2 is to compare (empirically) different choices of $S$, specifically to understand the trade-off between test PPL and estimated memory usage. The results suggest that Rademacher matrices offer better memory efficiency at a slight PPL penalty, and randomized semi-orthogonal matrices offer an even smaller penalty, though perhaps not a memory improvement (since it is not as clear how one could rematerialize them efficiently).

The proposed modification to training in the distributed setting definitely sounds clever [...] however, this seems to be specifically in the one-sided LoRA setting [...]

Yes, you are correct---this result is specific to the one-sided LoRA setting. We focused on this setting because this is what the theorems we prove focus on. That said, we also tried running a full two-sided distributed training setting (which would be equivalent to distributed ReLoRA/LTE; 20.97 PPL [200M] and 13.72 PPL [1.3B]) and found it to underperform the one-sided LoRA with worker-aware initialization. We will include these results in the next version of the paper.

Do you have an explanation for why the SVD curve in Figure 1 plateaus over time?

Great question! We suspect that towards the end of training, since the learning rate has annealed, the model is actually shifting very little in parameter space. This means that the gradient is also shifting fairly little. This means that our SVD-based gradient estimator, which is optimal (in the Frobenius error sense) when applied to the gradient the SVD was computed on, remains fairly optimal.

I was slightly confused by what the authors meant by the "Transformer++" architecture

Thanks for pointing this out. This part is unclear, and we will fix this in the next version of the paper. By Transformer++ all we mean is a Transformer with the changes made in the Llama architecture, e.g., RMSNorm instead of LayerNorm, SwiGLU activations, rotary embeddings (some papers use this term).

Other comments/suggestions

Thank you for your comments here—we will address them all in the next version of the paper. We also include a brief explanation/clarification where appropriate:

Please define the notation vec^-1 first used on line 195.

If $\text{vec}(\cdot)$ sends a matrix to its vectorized form (i.e., we stacked up the columns), then $\text{vec}^{-1}(\cdot)$ takes a big stacked vector and unstacks it into a matrix of the right size.

I didn't fully understand the points about why different schemes need different persisted matrices in Table 1

This point is related to your earlier question about Table 2. Table 1 is trying to be explicit about the cases where it should be possible to just persist the seed that generates a matrix, and regenerate it on-the-fly as needed. This means additional memory is saved since we can get away with just storing the seed for that matrix (~4 bytes), as opposed to storing the matrix itself.

I also found the discussion around Table 4 slightly confusing. It seems to show that in the baseline, having more workers is beneficial, but in the lora experiments, having more workers is detrimental?

Great question---In Table 4 we simultaneously vary the number of workers and the rank of the gradient transformation/LoRA adapter, except in the case of the baseline, where we only vary the number of workers. Since we allow all workers to train for the same number of tokens, this means that for the baseline, increasing the number of workers increases the effective number of total tokens we train on without any penalty (i.e., reducing the rank). So the baseline will always benefit from more workers without any real negative effects, whereas the other approaches will suffer from the rank being decreased as the number of workers is increased. We mention this briefly in the caption, but we will make sure to make this clearer.

It would be useful to clarify what the metrics are in all table captions (I think they are all test perplexity?).

That’s right, they are all test perplexity.

Thank you for your response! We will break down your question into parts:

Could you further explain your findings from experiments in Sec 4.1?

The key goal of Sec 4.1 is to understand the interplay between the choice of gradient transformation S and overall performance. One important aspect of this interplay is that some transformations are expensive to update (e.g., the SVD transformation works well according to the GaLore paper, but requires running an SVD which is expensive), and others take more memory to store (e.g., the SVD transformation requires you to persist that matrix in memory, but if you are using random Gaussian matrices you just need to store a seed that you will use to regenerate this matrix—this means the whole matrix can be compressed into a handful of bytes).

Table 1 tries to summarize these trade-offs conceptually. For each gradient transformation method, we describe what the parameters (of linear layers) in the model would look like using this approach, which of these parameters would be trained/frozen, and which of these need to be persisted in memory.

Table 2 shows the empirical results (in terms of PPL on a held-out subset of SlimPajama) of training 200M and 1B using each of these methods; since the LoRA formulation is amenable to quantizing the frozen weight, we also include for each method the results using INT8 quantization and NF4 quantization. Beyond the test PPL, we also report the estimated memory usage of each approach.

The most important point of Table 2 is that, at a minor loss of PPL (i.e., +0.1-0.3 PPL), one can use other transformations that offer memory efficiency (e.g., Rademacher gradient transformations). The other takeaways from the table are that (i) while two-sided gradient transformations may be intuitively nice, they don’t seem to work as well, and (ii) INT8 quantization gives you memory savings without any meaningful PPL degradation, but NF4 starts to incur a non-trivial PPL penalty.

But we agree that section 4.1 could be explained better and will make sure to expand on the above points in the next iteration of the paper.

The empirical study is not clear enough to me, regarding supporting the previous claim about the equivalence between weight-efficient optimization and gradient-efficient optimization. Table 2 is more like a performance comparison instead of showing any connections between different methods [...] how do these results support the equivalence between weight-efficient optimization and gradient-efficient optimization

Since we prove that there is a mathematical equivalence between training with gradient transformations and training with a weight transformation (aka., a linear, one-sided adapter), the goal of our experiments is not to empirically verify this equivalence (since this would be redundant), but instead to leverage it by exploring, e.g., the interplay between choice of gradient transformation and empirical language modeling performance. That said, we will add a small scale experiment in the appendix where this equivalence is validated empirically, by showing that the loss curves of two toy models remain the same during training.

[...] and how is this equivalence used to build more efficient training methods?

Beyond the new gradient transformations we propose in Table 2, our biggest contribution in terms of more efficient training methods is in combining gradient transformations with distributed training. In particular, in Table 3 we show that making different workers have different gradient transformations and making these transformations pairwise orthogonal (i.e., the distributed random method) is better than having them be orthogonal (but not pairwise orthogonal) or having them be the same across workers. In Table 4, we also study the effect of varying the number of workers, and find that the benefits of the distributed random method become more pronounced as the model is partitioned across more and more workers (each training a smaller slice of the model).

Thank you for your review! We identify three key criticisms pointed out in your review, and address them below

1: Conceptual insights in the duality derivation are simple

Indeed, while the essence of our derivation was observed previously in the literature, our derivation is more general; it considers a linear adapter applied to any parameter vector and it takes into account optimizer choice (previous works were only showing the equivalence for parameter matrices of linear layers and SGD, a non-stateful optimizer). This view is productive: it allows us to establish a connection between Kronecker-factored gradient transformations and recent two-sided variants of LoRA adapters which have been found to perform well for parameter-efficient finetuning, in particular MoRA, PMSS, and LoRA-XS. This is to the best of our knowledge novel in the literature. We think that this more general result is a worthwhile contribution to the literature.

2: The empirical contributions borne out by the insights of the derivation are incremental.

A key criticism seems to be that despite proposing other gradient transformations (e.g., Rademacher, randomized semi-orthogonal), the SVD transformation (i.e., GaLore) is still the highest performing transformation in Table 2 in terms of PPL in the 1B case. However, in Table 2 we see that the SVD transformation underforms the proposed approaches in the more memory-efficient NF4 quantization (i.e., QLoRA) setting: the Rademacher and Random-Orthogonal projections in this setting perform nontrivially better. Furthermore, in the 1B setting without quantization (or with INT8 quantization), the gap between the SVD and some of the other approaches is fairly modest (e.g., 0.1-0.3 PPL), and it is important to take into account that some of these are more memory efficient than the SVD transformation (see memory column in Table 2 for, e.g., Rademacher transformations).

We also kindly note that a large portion of our paper was dedicated to the application of the gradient transformation–adapter duality to distributed training, though it seems this wasn’t mentioned in the review. In particular, we show that choosing the transformations so they are worker-aware leads to better results than making them randomly different across workers or making them the same across workers.

Finally, we want to emphasize that the goal of the paper wasn’t necessarily to achieve state-of-the-art memory-efficient pretraining results, but more to give a flavor for how our duality can unify existing approaches and lead to new techniques.

3: Discrepancy between our results and those in the GaLore paper

First, we’d like to stress that our setup is different from the GaLore paper in many respects, so it is hard to do a direct comparison. For example, we use a different dataset (SlimPajama instead of C4), we use gradient accumulation, our maximum sequence lengths are substantially higher (2K instead of 256), and our optimizer hyperparameters are not the default Adam hyperparameters, but instead are more consistent with values used in more recent, large LLM training runs (we use the ones for Llama-2 and OLMO).

More importantly, our decision to use a different setup was deliberate: Since we focus on pretraining, we want to be in the setting that most closely resembles recent LLM training runs. Therefore, we do not see the discrepancies between our results and those in the GaLore paper as a reflection that our setup is suboptimal; rather, we see our results, in part, as an attempt to apply GaLore to a different setting. We also note that such nontrivial differences in performances due to different training setups is common: for example, the Flora paper’s replication of GaLore (in Table 6 of the appendix) also shows substantially different results.

Answers to individual questions

1. See point 3.

2. This is a great question—we were also puzzled by this. Here, we actually think it is worthwhile taking a step back and looking at the results and models more broadly. Mathematically (from our duality derivation), it is clear that GaLore (and thus one-sided LoRA training) partitions the parameter space and trains only along a subspace of it. This reduces the effective number of parameters in our model, which would usually lead to a performance hit. Further, from our derivation, ReLoRA can be seen as GaLore where the projection is learnt by the optimizer, so unless this small number of new parameters introduce a lot more instability, we arguably shouldn’t expect the results to change much. Both of these observations manifest in our results. From this lens, we think our results are arguably what we’d a priori expect. While we cannot be sure why previous work doesn’t also show this same trend, but we suspect that this may have to do with the differences in the experimental setup (see point 3 above).

3. See point 2.

Thank you for your review!

Our main objective in this paper was to explore pretraining, and thus we selected a pretraining dataset that matches the distribution of pretraining data (SlimPajama, which is based on RedPajama, that tried to match the pretraining data used by Llama, one of the most successful open-source LLMs).

In terms of model size, because of compute constraints in an academic setting, we scaled as large as we could go (1B). As much as we’d love to scale further, we cannot do a pretraining run of a larger scale on our resources.

I wrote more comments above but how the convergence or the collapse of the method will be with this method?

Besides the standard convergence results that apply to ML models of this scale, we note that, due to the nature of our equivalence, the theorems in the GaLore paper should still apply (subject to those theorems’ conditions, e.g., any restrictions on the choice of projection). Empirically, we find that all approaches we consider have similar convergence behavior.