ScaleBiO: Scalable Bilevel Optimization for LLM Data Reweighting
Authors: Rui Pan, Jipeng Zhang, Xingyuan Pan, Xiaoyu Wang, Renjie Pi, Tong Zhang


Reviewer paper summaries: This paper focuses on applying bilevel optimization to reweight data sources in training/fine-tuning LLMs. The high-level idea is to transform bilevel optimization (Eqn.1) to a minmax optimization problem (Eqn.5). The authors propose the ScaleBiO algorithm to solve such minmax problem with some auxiliary variables. For each training step, we will further sample a block from the models to perform the alternative update (similar to the LISA algorithm). The convergence rate is $O(1/K^{2/7})$ that matches Kwon et al. The authors evaluate the performance on whether the algorithm can successfully determine the correct data weights in multiple small scale experiments (data denoising, multilingual weighting, etc.), and for large-scale experiments, the authors evaluate on the final MT-bench scores for LLaMA3-8B (and transferability on LLaMA3-70B).
Summary of strengths: Table 1 is quite informative on positioning ScaleBiO between the related works.

The transformation from bilevel optimization to minmax optimization is clear. The motivation behind applying bilevel optimization to data reweighting is clear. The overall paper flow on section 3 is clear.

Multiple applications (data denoising, multilingual reweighting, differentiating high-quality from low-quality dataset) from data reweighting are considered and evaluated. 

Theoretical convergence rate of ScaleBiO is provided, although it is better to have a comparison similar to Table 1 with the results from the related works, especially given that the minmax formulation of bilevel optimization is *not* new.
Summary of weaknesses: 1. Lack of proper experiment baselines and complete assessment

Baselines:

The small-scale experiments all (Figure 1, 2, 3, 4) don't have a baseline. More importantly, the baseline in Table 3 is not sufficient. We would need a baseline that also uses bilevel optimization methods (1st order) for a more complete and fair comparison. 

Complete assessment:

Since ScaleBiO uses blockwise update for memory efficiency, we should profile the memory consumption across LLMs with multiple sizes and number of blocks. On the optimization side, the effect of smoothing constant $\alpha$ should be adequately evaluated.

2. Comparison with related works on the theoretical properties

We should note that the minmax formulation of bilevel optimization is not new. As ScaleBiO is fairly simple on the optimization side, a paragraph of comparison on the theoretical properties of ScaleBiO should be included in the paper.
Comments, suggestions and typos: This might be a typo or I am misunderstanding the algorithmic procedure of ScaleBiO:

Why do we have a gradient descent step for auxiliary variable $u$ (Line 5, Algorithm 1) instead of a gradient ascent step for $u$ as we are maximizing $L_\alpha$ w.r.t. $u$ ?

Line 953 has an empty space between the bracket ` (see Lemma [])`.
Overall assessment scores: 2.5