Synthesize, Partition, then Adapt: Eliciting Diverse Samples from Foundation Models

Thank you for your insightful review. We'd like to address your main concerns.

Parallel Generation and Memory Efficiency

Recent advances in serving multiple LoRA adaptations in parallel significantly mitigate the need to store multiple full checkpoints. Notably, the S-LoRA system [1] demonstrates that it's possible to efficiently serve thousands of LoRA adapters with minimal overhead, using techniques like Unified Paging to manage memory efficiently. Furthermore, FLoRA [2] allows each input example in a minibatch to be associated with unique low-rank adaptation weights, enabling efficient batching of multiple LoRA adapters. In these systems, we only need to store the LoRA weights in memory, which are often low-rank, significantly reducing memory requirements. These advancements suggest that SPA can be implemented without significant memory constraints or parallel generation difficulties. As the field continues to progress, we anticipate even more efficient solutions for managing multiple LoRA adapters.

[1]: Sheng, Y., Cao, S., Li, D., Hooper, C., Lee, N., Yang, S., Chou, C., Zhu, B., Zheng, L., Keutzer, K., Gonzalez, J.E., & Stoica, I. (2023). S-LoRA: Serving Thousands of Concurrent LoRA Adapters. MLSys 2024

[2]: Wen, Y., & Chaudhuri, S. (2023). Batched Low-Rank Adaptation of Foundation Models. ICLR 2024

Best number of partitions

Determining the optimal number of partitions is indeed an important consideration. This highly depends on the size of the synthetic dataset, with larger datasets typically inducing a higher optimal number of partitions. We conducted an ablation study (Fig. 6 in our paper) examining the effect of the number of partitions on diversity. We found that the diversity score remains relatively stable as the number of adaptations increases from 8 to 12, suggesting that 8 partitions is roughly the sweet spot for the synthetic dataset we used (75k data points). We also observed improvements in diversity when increasing the number of partitions from 4 to 8, though this data was not included in Fig. 6.

Computational time

We provided a rough complexity in the limitation section. To offer more details: One epoch of fine-tuning on the 75k synthetic data points takes approximately 8 hours using 2 A40 GPUs. We trained the LoRA weights for 4 epochs. Influence function calculation costs roughly 1 epoch of the training time, which is approximately 8 hours. We want to highlight that the main computational cost occurs during the offline partitioning and adaptation phases. Future work could explore more efficient data attribution methods, such as TRAK [1] and K-FAC [2], which have the potential to substantially reduce this overhead while maintaining the benefits of our approach.

We appreciate your feedback and hope our clarifications have addressed your concerns about parallel generation and memory requirement. If you have any remaining concerns, please let us know. We're grateful for the opportunity to further explain our work.

[1]: Park, S.M., Georgiev, K., Ilyas, A., Leclerc, G., & Madry, A. (2023). TRAK: Attributing Model Behavior at Scale. International Conference on Machine Learning. ICML 2023

[2]: Grosse, R.B., Bae, J., Anil, C., Elhage, N., Tamkin, A., Tajdini, A., Steiner, B., Li, D., Durmus, E., Perez, E., Hubinger, E., Lukovsiut.e, K., Nguyen, K., Joseph, N., McCandlish, S., Kaplan, J., & Bowman, S. (2023). Studying Large Language Model Generalization with Influence Functions.

Thank you! I am afraid that I would still keep my score because in a high level, I am not quite in favor of the idea of having multiple lora weights.

Thank you for your follow-up question. We conducted two simplified scenarios using S-LoRA codebase and vLLM to compare sampling from multiple LoRA components versus a single model. Our setup uses a synthetic dataset in S-LoRA repo with 600 prompts and responses ranging from 8 to 512 tokens respectively. For each prompt, the model generates tokens until the output length matched the response's length. We used Llama-13b with LoRA rank 16, leading to a 10% memory overhead for storing 8 LoRA adapters compared to a single model. When running the serving experiment, the requests are coming with a request rate of 2 per second.

In the first experiment, we simulated a scenario where multiple GPUs were used for serving. Here, the goal was to generate 8 diverse samples per request. We used 8 A100-40G GPUs. For the single model setup, incoming requests (say x and y) were expanded into batches of 8 identical prompts, $[x_1, x_2, ... , x_8]$ and $[y_1,y_2,...,y_8]$ and handled by vLLM’s continuous batching and routing. Usually, it will distribute $[x_1, x_2]$ to GPU-1 job queue, $[x_3, x_4]$ to GPU-2 job queue and so on.

For the SPA model, the batch becomes $[x_i, y_i, LoRA_i]$, where i ranges from 1 to 8. The router then distributed these to different GPU job queues, with each GPU handling a unique prompt-LoRA pair. This also gives us 8 samples per request. The results showed that this leads to a minor 3% overhead in both throughput (tokens / sec) and latency (secs / output token), primarily due to the additional step of distributing the LoRA adapters. This approach ensures parallel generation but requires multiple GPUs. However, it is a common practice to use multiple GPUs to serve real-world requests.

In the second experiment, we explored a single GPU scenario using S-LoRA. This setup also stores all LoRA weights in memory, leading to a similar 10% memory overhead as in the first scenario. In this experiment, we only generated one sample per request for simplicity. This means that for the single model, the metrics should be roughly the same as those obtained when sampling 8 outputs using 8 GPUs in the above scenario. In the SPA setup, each request was coupled with a random LoRA adapter, which is handled by the S-Lora library. The result below shows that this configuration leads to a ~10% overhead in throughput and latency.

The first result indicates that SPA can be implemented efficiently in a multi-GPU setup with almost no performance lost. Even in a single GPU scenario, leveraging S-LoRA allows us to sample from multiple LoRA adapters without significant trade-off. This is also an active research field, we could anticipate even more efficient solutions for sampling from multiple LoRA adapters with some future works. We hope these results help address your concerns.

Thank you for your insightful review. We'd like to address your main concerns.

Accuracy on NL tasks

The accuracy of SPA methods (influence-based, lexical-based, and random) all achieve similar accuracy to the single model, reaching the same conclusion as the pass@1 metric in Fig. 4a. We didn't observe improved pass@k with increased diversity on these tasks, likely because most are multiple-choice questions. For diversity measurement, we asked models to continue generating text even after producing an answer choice. We omitted accuracy results as they didn't provide new insights beyond the code generation tasks. However, we acknowledge that including this data would have provided a more complete picture and will add it to the revised paper.

Trade-off in high temperature sampling, temperature plot stops at 0.5

You're correct that at higher temperatures, the diversity of single models converges to SPA models. In practice, low temperature or even greedy sampling is often preferred when average sample quality is crucial. As shown in Fig. 4, high temperature leads to a decrease in the pass@1 metric, indicating a drop in sample quality. High temperature sampling may deviate from the learned distribution and produce hallucination or less coherent outputs [1]. We briefly discussed the disadvantages of high temperature in the introduction and will expand it in the revision.

We limited our plots to temperatures up to 0.5 as SPA's primary goal is to maintain average sample quality (measured by pass@1) while achieving diversity. Fig. 4a shows that temperature 0.5 already leads to a significant drop in pass@1. However, we acknowledge the importance of understanding where methods converge in terms of diversity. Our preliminary experiments indicate convergence around temperature 0.8. We will include these higher temperature results in our revision.

Number of adapters

The optimal number of adapters highly depends on the size of the synthetic dataset, with larger datasets typically inducing a higher optimal number of adapters. We observed improvements in diversity when increasing the number of adapters from 4 to 8, though this data was not included in Fig. 6, suggesting that 8 adapters is roughly the sweet spot for the synthetic dataset we used (75k data points).

KL divergence, hyper-parameters, distribution over LoRA adapters

We measured the KL divergence at the first decoding step, which is at the token level. This measures the difference in states across different adaptations after seeing the same prompts. This is not applicable to the single model approach.

For each method (when temperature > 0), we sampled a total of 120 samples to calculate pass@1, pass@k, and diversity score (K=5). This larger sample size helps reduce metric variance. For SPA methods, we distributed samples evenly among adapters (15 samples per adapter). Future work could explore weighted sampling in SPA. The exception is greedy sampling (temp=0 in Fig. 4 and Sec. 5.3), where we only get one sample per adapter, leading to a total of 8 samples.

Why pass@1 is not higher, how to select queries.

SPA is primarily designed to improve diversity, which doesn't necessarily improve average sample quality. Hence, it shows similar pass@1 performance to the single model.

Our principle was to create a diverse query set. We used few-shot prompting with ChatGPT to generate ~20 queries, then manually selected 8 covering a wide range of topics.

Single model at higher temp in Sec. 5.3

We didn't include comparisons with single models at higher temperature because, as shown in Fig. 4a, higher temperatures lead to sample quality drops. However, we acknowledge that including this comparison could provide valuable insights. In our revision, we will add this comparison and discuss the trade-offs.

We appreciate your feedback and hope this clarifies your concern. Please let us know if there is any remaining concerns.

[1]: Lee, M. (2023). A Mathematical Investigation of Hallucination and Creativity in GPT Models. Mathematics.

Thank you for your positive feedback. Here are the accuracy results:

SPA maintains comparable accuracy to the single model. The GPQA evaluation has high variance; even small changes in the prompt can lead to big differences in accuracy. Despite this, we still included GPQA because it provides a good source for evaluating diversity due to its broad range of topics. We will include the distribution in the next revision as you suggested.

We've also added new experiments in our reply to reviewer 4Daa addressing concerns about parallel generation and memory overhead. We greatly appreciate your thoughtful feedback and hope we have addressed your concerns. We're keen on your support for our submission.

Thank you for your insightful review. We'd like to address your main concerns.

Motivation of why doing SPA and why doing so leads to diverse samples. Role of synthetic data and the generalizability of it.

We would first like to clarify that LLMs must typically go through synthetic data tuning before deployment, whether for domain adaptation, value alignment, or distilling knowledge from a stronger model. SPA is designed to utilize this existing synthetic data used in LLM post-training stages (or mid-training stage). Operating within this existing paradigm, SPA addresses the inefficiency of tuning a single model on increasingly larger synthetic datasets due to diminishing returns.

The first motivation, as stated in Section 3, is to turn this potential inefficiency into an advantage for diversity from the multi-model perspective. The second motivation is sampling from multiple distinct distributions (i.e., different model adaptations) naturally yields more diverse outputs compared to sampling multiple times from a single distribution (i.e., a single model). This is why SPA leads to diverse samples by creating multiple model adaptations, each trained on a different partition of the synthetic data. Importantly, due to diminishing returns, training on data partitions doesn't significantly affect accuracy, as shown in Sec. 3.

Furthermore, our partitioning strategy, especially when using influence functions, creates subsets of data that induce distinct model skills. This further enhances the diversity of the resulting adaptations (compared to the random and lexical partitions).

Regarding generalizability, since SPA uses data already part of the post-training process, it doesn’t introduce new generalization risks. While our method benefits from more diverse data, it's designed to improve diversity even when working with less diverse synthetic datasets, by being able to sample from multiple distinct distributions.

More comparisons and related works

Thank you for this feedback. For the additional comparisons, please see the general response. We didn't compare to QD search methods because SPA and QD methods operate on fundamentally different principles. SPA operates at the model level, creating multiple models during an offline process. The main computational cost occurs during the offline partitioning and adaptation phases. This allows for quick, parallel diverse sampling at inference time. In contrast, QD search operates at the generation level, iteratively going through mutation, evaluation, and refinement cycles. It introduces overhead during sampling, making it less suitable for real-time applications. We will also expand our literature review to include all the suggested works and other relevant studies.

Connection between partitioning and diverse LoRA components

The influence function captures how training examples affect model predictions on specific test queries. Our partitioning strategy groups examples with similar influence patterns, resulting in K clusters. Each cluster consists of synthetic data that most strongly influences a particular test query (detailed at line 257). When fine-tuning LoRA on these individual clusters, each component focuses on the skills emphasized by its associated test query.

Importantly, we don't require extremely diverse LoRA components because the diversity primarily stems from sampling across multiple distinct distributions (motivation #2). Even with random partitioning, this strategy ensures that these distributions are sufficiently different to generate diverse outputs when sampled collectively.

Generating more than K diverse responses

While we train K adapters, SPA is not limited to generating only K diverse responses. Each adapter can produce multiple outputs, especially when combined with other sampling techniques (e.g., temperature sampling). In our experiments, for temperatures > 0, we generated a total of 120 samples to calculate pass@1, pass@k, and diversity scores. For SPA, this meant sampling 15 samples from each of the 8 adapters. In contrast, for the single model baseline, all 120 samples were drawn from the same distribution. This sampling strategy highlights a key advantage of SPA: even when generating many more than K samples, we're still drawing from multiple distinct distributions. This inherently promotes diversity compared to repeatedly sampling from a single distribution. Moreover, SPA is complementary to other diversity-enhancing methods.

Quality and bias in randomly sampled LoRA components

Our experiments show that random sampling of LoRA components during inference doesn't significantly affect response quality, as evidenced by our pass@1 results (also because of motivation # 1). Our partitioning strategy aims to create balanced subsets that cover different aspects. The additional comparisons in the general response showed that SPA is less biased than diversity of thought. However, we acknowledge that some components might specialize more in certain areas. In practice, a held-out dataset can be used to identify and remove overly biased components before deployment. Future research could explore weighted component selection, using a similar routing strategy in MOE.

Selection of M queries

The M questions don't necessarily need to be manually chosen. In our experiments, we used a semi-automated approach: we few-shot prompted ChatGPT to generate ~20 diverse queries and then manually selected 8 of them to cover a wide range of topics (Appendix C).

We appreciate your feedback and hope our clarifications have addressed your concerns about motivation and the empirical results. If you have any remaining concerns, please let us know. We're grateful for the opportunity to further explain our work.

We sincerely appreciate your reconsideration and the increased score. We'd like to take the opportunity to address your remaining concerns:

The same training procedure for each LoRA component leads to less diversity

Our approach operates under the assumption that LLMs undergo post-training using synthetic data, which is a common practice (e.g., iterative post-training in Llama 3 utilizes synthetic data in each round). Under this assumption, we have limited flexibility to modify the training procedure. Given these constraints, and motivated by diminishing returns, sampling from multiple distributions using LoRA components is our strategy to outperform a single model tuned on the entire synthetic dataset. While each LoRA component on its own may face diversity challenges like you mentioned, it won't be worse than the single model scenario.

Moreover, if modifying the training procedure is allowed, we could incorporate the mutual information as one of the objectives when optimizing LoRAs, this will encourage these LoRA components to be as diverse as possible.

Quality-diversity trade-off

We agree that calibrating this trade-off may be challenging in practice. On the other hand, SPA offers an improved Pareto frontier, improving diversity even at low temperatures. Users can now achieve good diversity at low temperatures (even when greedy), which was previously difficult. Thus, while SPA doesn't eliminate the trade-off, it provides users with more options.

Predefined K and component quality, get the most probable sample

We want to clarify that K is not arbitrarily selected but depends on the synthetic dataset size and the strength of diminishing returns. As shown in Fig 4a, $\sum_i^K \text{Accuracy}(\text{model}_{\phi_i}) / K$, roughly equals the accuracy of a single model $\phi$, fine-tuned on the entire synthetic dataset. The average quality of LoRA components is roughly the same as the single model baseline.

For cases requiring the most probable answer (only one sample is required), we can conduct an offline evaluation on a held-out dataset before deployment, then sample from the best-performing LoRA component.

We hope our clarifications can further address your concerns.

Thank you for your insightful review. We'd like to address your main concerns.

In-context v.s. Fine-tuning

We can formulate the difference as sampling from $P_{\phi_1} (y|x)$ and $P_{\phi_1}(y|x')$ for in-context versus $P_{\phi_1}(y|x)$ and $P_{\phi_2}(y|x)$ for our approach ($\phi$ denotes model parameters). We argue that sampling from multiple distinct distributions (SPA) produces more diversity than modifying the prompt. In the general response, we conducted additional experiments comparing SPA with Diversity of Thought, an in-context approach. The results show SPA's significant advantage. Moreover, one could potentially combine them by sampling from $P_{\phi_1}(y|x)$ and $P_{\phi_2}(y|x')$, further enhancing diversity. Finally, we assume the LLM must go through the synthetic data tuning before deployment, whether for domain adaptation, value alignment, or distilling knowledge from a stronger model. Given the diminishing returns observed with increasing synthetic data, the motivation for SPA remains strong.

Role of LoRA rank

We agree that your intuition about LoRA rank is correct. A smaller rank would lead to more similar representations, potentially reducing diversity, while a larger rank might enhance it. We used a fixed rank for consistency in our experiments, but exploring different ranks is an interesting direction for future work. Yet, the benefit of sampling from multiple distinct distributions remains valid regardless of the rank.

Number of adapters

The optimal number of adapters highly depends on the size of the synthetic dataset, with larger datasets typically inducing a higher optimal number of adapters. We observed improvements in diversity when increasing the number of adapters from 4 to 8, though this data was not included in Fig. 6, suggesting that 8 adapters is roughly the sweet spot for the synthetic dataset we used (75k data points).

Further improvements

Thank you for your value suggestions for new diversity measures and how to improve the intro section. Regarding the new diversity measure, we'd appreciate clarification on training a new model on generated datasets, as our single model adaptation baseline is already trained on the entire synthetic data. We’ll summarize the method more concisely in the intro and move Fig. 3 ahead in our revision.

We appreciate your positive feedback and hope this addresses your concerns. We kindly ask for your support of our paper.

Thank you for your positive feedback. We've also added new experiments in our reply to reviewer 4Daa addressing concerns about parallel generation and memory overhead. We're keen on your support for our submission.