Set-LLM: A Permutation-Invariant LLM

We sincerely thank Reviewer f4vF for the careful review and the thought-provoking questions that help us to improve our paper.

W1: The paper targets a niche aspect of LLMs and may have little to no practical use.

Respectfully, we do not agree with the reviewer that this is a niche aspect of LLMs with little to no practical use. Many works identify order bias as a problem for LLMs, and we present an effective, theoretically grounded solution. While we chose Multiple-choice QA for benchmarking, we emphasize that our method extends to other LLM tasks, such as LLM as a judge, multi-document summarization, and retrieval-augmented generation. To strengthen our paper, we have conducted additional experiments on these tasks. The results validate the general applicability and benefits of the proposed method. We will expand these and include them in the final paper.

Multi-Document Summarization

We conduct our evaluation on the MultiNews [1] dataset. We filter out inputs of length greater than 20000 characters and use the standard train-validation-test split. We compare the original model architecture with Set-LLM with single inference runs (and no adversarial setting). We use the best-performing hyperparameters from the multiple-choice experiments.

We report standard metrics: Rouge F1 scores and compression (the length of the original text divided by the length of the summary). We also include Rouge F1 scores between the target summaries and the model summaries (R1*, R2*, R3*).

Set-LLM outperforms the finetuned baseline with all three base models in all metrics. However, the margin is significantly narrower with the larger Gemma 7B model. Note that the compression rate is always lower for Set-LLM, meaning the summaries it produces are longer.

[1] Fabbri, Alexander R., et al. "Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model." arXiv preprint arXiv:1906.01749 (2019).

Multi-Document Question Answering (Retrieval-Augmented Generation)

Motivated by [2], we evaluate Set-LLM on retrieval-augmented question answering. [2] introduces a task mimicking the RAG setup. They ask questions with 20 supporting documents, where only one document contains the information required for the answer. They show that LLMs (e.g., GPT3.5 Turbo) produce a U-shaped performance curve given the location of the answer. That is, if the document with the answer is listed first or last, then the LLM answers the question correctly with a much higher probability (e.g., 75% in position 1/20 vs 55% in position 10/20). In contrast, Set-LLM is guaranteed to have a flat performance curve, since reordering the documents does not affect the outcome.

We split the dataset into 60% training, 10% validation, and 30% test data. We replicate the training data 5 times and shuffle the documents randomly for each input. This is done to avoid the model overfitting to a specific formulation of each training question. Again, hyperparameters were taken from the multiple-choice experiments.

We have 5 different test sets with the answer document placed in the first, 5th, 10th, 15th, and last positions. We provide results for two LLama-based models below. All scores denote percentage accuracy.

[2] Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2024. Lost in the Middle: How Language Models Use Long Contexts. Transactions of the Association for Computational Linguistics, 12:157–173.

We sincerely thank Reviewer nn3i for the positive feedback and thoughtful suggestions.

W1: Requires finetuning. That said, the computational overhead seems manageable.

We agree with the reviewer, that the proposed fine-tuning method results in an additional training step, however the amount of fine-tuning (4h on a single GPU) is negligible in comparison to the overall pre-training time (e.g. 1.46M GPU hours for llama 3.1 8B) and the resulting model can be used in a range of applications.

W2: It would strengthen the paper to verify that general model capabilities do not degrade on standard benchmarks.

We endeavour to provide an evaluation of the models after fine tuning on standard evaluation benchmarks they were built for during the author reviewer discussion period. We have machine translation and math in mind, but would be happy to take suggestions. Note, that in response to other reviewers, we have evaluated Set-LLM on additional tasks (multi-document summarization and retrieval-augmented generation).

Q1: Why change the causal nature of the mask?

This is an interesting suggestion. The overall causal structure of the mask has to be destroyed in order to guarantee invariance. However, the attention within an option could remain causal. We have not tried this, but it would be an interesting addition to our ablation studies. We do not expect significant performance gains since a fully-connect prefix mask is shown to be at least as effective as a causal mask in [1].

[1] Kopiczko, Dawid J., Tijmen Blankevoort, and Yuki M. Asano. "Bitune: Bidirectional instruction-tuning." arXiv preprint arXiv:2405.14862 (2024).

Q2: There seems to be a missing notation in the statement of Theorem 1, what is \tilde{X}(t+1) ?

Thank you for bringing this to our attention. \tilde denotes variables after permuting the input. We will make sure to add this clarification.

Q3: In Table 2, it is surprising that the PrefixMask + No positional embeddings configuration still yields meaningful outputs.

Indeed, this is very surprising, that a bag-of-words, or in this case better a bag-of-tokens model performs so well, especially the 74.54% on PIQA. This probably says more about the dataset than about the model. Note that the performance on the other datasets is weaker.

Q4: Concerning the adversarial order: do the authors generate multiple completions per input order, or just one?

We do not generate completions as answers, instead, the possible options are ranked by their log probabilities. Therefore the evaluation is deterministic once the model has been trained.

We sincerely thank Reviewer ryD4 for their feedback and for helping us to improve our paper.

W1: The scope for this work seems to be rather limited.

We understand the concern. We argue in the paper that the multiple choice setting is a reasonable test case for LLM-as-a-Judge applications, and LLM-as-a-Judge has wide-ranging applications. We mention other possible applications in the paper and see our model as a possible blueprint for these applications. To strengthen our paper and showcase the versatility of our approach we have conducted additional experiments on multi document summarization and retrieval augmented generation, both common LLM tasks. The results validate the general applicability and benefits of the proposed method. We have not been able to run all the baseline models on both of these applications, but we will make sure to complete these results for the final version of the paper.

To strengthen our paper and showcase the versatility of our approach we have conducted additional experiments on multi document summarization and retrieval augmented generation, both common LLM tasks. The results validate the general applicability and benefits of the proposed method. We have not been able to run all the baseline models on both of these applications, but we will make sure to complete these results for the final version of the paper.

Multi Document Summarization

We conduct our evaluation on the MultiNews [1] dataset. We filter out inputs of length greater than 20000 characters to satisfy our memory constraints. We use the standard train-validation-test split. We compare the original model architecture with Set-LLM with single runs (no adversarial setting). Our initial results do not include hyperparameter tuning, we use the best-performing hyperparameters from the multiple-choice experiments. The experiments will be expanded for the final paper.

We report standard metrics, this includes Rouge F1 scores and compression - the length of the original text divided by the length of the summary. We also include Rouge F1 scores between the target summaries and the model summaries (R1*, R2*, R3*).

Set-LLM outperforms the finetuned baseline with all three base models in all metrics. However, the margin is significantly narrower with the larger Gemma 7B model. Note that the compression rate is always lower for Set-LLM, meaning the summaries it produces are longer.

[1] Fabbri, Alexander R., et al. "Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model." arXiv preprint arXiv:1906.01749 (2019).

Multi-Document Question Answering (Retrieval-Augmented Generation)

Motivated by [2], we evaluate Set-LLM on retrieval-augmented question answering. The authors introduce a task mimicking the retrieval-augmented generation setup. They ask questions with 20 supporting documents, where only one document contains the information required for the answer. They show that commercial LLMs (GPT3.5 Turbo) produce a U-shaped performance graph given the location of the answer. If the document with the answer is listed first or last, then the LLM answers the question correctly with a much higher probability (e.g. 75% in position 1/20 vs 55% in position 10/20). In contrast, Set-LLM is guaranteed to have a flat performance curve, since reordering the documents does not affect the outcome.

We split the dataset into 60% training, 10% validation, and 30% test data. We replicate the training data 5 times and shuffle the documents randomly for each input. This is done to avoid the model overfitting to a specific formulation of each training question. We have not conducted hyperparameter tuning yet, the hyperparameters were taken from the multiple-choice experiments.

We have 5 different test sets with the answer document placed in 5 different locations (first, 5th, 10th, 15th, last). We provide the results for two LLama-based models below. All scores denote the percentage accuracy.

[2] Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2024. Lost in the Middle: How Language Models Use Long Contexts. Transactions of the Association for Computational Linguistics, 12:157–173.

W2: I would question if this is necessary to finetune a model just to make it work for multi-choice questions.

To add to the response to W1, the finetuning cost is fairly small (between 1 and 4 GPU hours) for each dataset and model we consider. This is negligible in comparison to the millions of GPU hours used to pretrain these models (e.g. 1.46M GPU hours for llama 3.1 8B). Moreover, the main purpose of finetuning is to allow the model to adapt to its new, altered architecture. We envisage that a typical scenario would involve finetuning Set-LLM on a general-purpose set-dataset and then using it for relevant downstream applications.

I am interesting in knowing the influence of the set-llm finetuning to non-multiple choice questions.

This is a great suggestion! It took some time to get results.

We chose two standard LLM benchmarks to test the effect of set-LLM finetuning on non-multiple-choice questions. For question-answering, we use 1000 test questions from SQuAD v2 [1], and for machine translation, we use the English-German subset from Flores-200 [2]. For SQuAD v2 we record accuracy, and for Flores-200, we record COMET scores.

As finetuned models, we picked models finetuned on PIQA (arbitrarily) and Long Context QA [3] (from the rebuttal). We test the models in 0-shot, 1-shot, and finetuned settings. We compare the original models with the Set-LLM models. PIQA models were finetuned on multiple-choice questions and therefore trained to copy a choice from the input to the output. Long QA models were finetuned to answer a question based on a set of (retrieved) documents.

SQuAD v2 (PIQA)

*Unanswerable questions were removed for 0-shot and 1-shot settings

SQuAD v2 (LongQA)

*Unanswerable questions were removed for 0-shot and 1-shot settings

Flores-200, English -> German (PIQA)

Flores-200, English -> German (LongQA)

The PIQA models were finetuned to copy a choice from the input to the output. We find that these models struggle with SQuAD in the 0-shot and 1-shot (to a lesser extent) settings. The LongQA models perform better. Gemma 2B Set-LLM, in particular, performs very poorly. For Flores-200, the differences are less pronounced.

Overall, we find that Gemma 2B Set-LLM underperforms Gemma 2B Original in some settings, but not when finetuned.

However, the Llama-based models generally do not have a significant gap between original and Set-LLM performance. This is true for both datasets and across almost all evaluation settings. The exception is Llama 3.2 3B finetuned on PIQA in 0-shot and 1-shot settings.

[1] Rajpurkar, Pranav, Robin Jia, and Percy Liang. "Know what you don't know: Unanswerable questions for SQuAD." arXiv preprint arXiv:1806.03822 (2018).

[2] Costa-Jussà, Marta R., et al. "No language left behind: Scaling human-centered machine translation." arXiv preprint arXiv:2207.04672 (2022).

[3] Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2024. Lost in the Middle: How Language Models Use Long Contexts. Transactions of the Association for Computational Linguistics, 12:157–173.

Why not evaluate on benchmarks like search arena where the evaluation is on LLM-as-a-judge?

The main aim of this paper is to create a set-invariant LLM architecture using pretrained LLMs. This then has several potential applications as discussed in the paper and the rebuttal. We believe that overall, the experiments convincingly show the viability of the approach. The aim of the paper is not to explore in detail the specific application of LLM as a Judge, though admittedly this may not have been clear in the initial submission, especially with the focus of the experiments on a single task. However, with the help of the feedback from the reviewers and the additional results, we have expanded the scope of the experimental section. As such, we feel a deep-dive into Set-LLM for LLM as a Judge (or other applications) should be deferred to future work. We have run some initial experiments on the suggested search arena dataset, but we have used the small pretrained LLMs from the paper, and the results are not conclusive.

Search Arena (finetuned)

*As suggested by Reviewer f4vF

We restricted the search arena dataset to questions with a single turn (not multi-turn), with a winner (no ties), and no longer than 5000 characters. We use a 2/3-1/3 train-test split.

We sincerely thank Reviewer Zm9C for the positive feedback and valuable suggestions to improve the paper.

W1: Some technical sections (e.g., the formalization of set-token mappings in Lines 104–107) are notation-heavy and complex and may benefit from clearer exposition or illustrative examples.

Thank you for raising this concern. We have tried to include figures and illustrative examples to help the reader, but we will certainly go through the paper again and see if we can improve clarity. Indeed, the set-token mappings sound much more complicated than they are, and we included Figure 3 for this reason. We will take a close look at whether we can simplify this formalization for the final version. We hope the intuition however is clear.

W2: While the multiple-choice QA setting is well-justified, experiments focus solely on this task; broader generalization to other set-input scenarios (e.g., retrieval-augmented generation, ranking) is left untested.

We performed additional experiments on retrieval-augmented generation and multi-document summarization. We describe these and discuss the results below.

Multi Document Summarization

We conduct our evaluation on the MultiNews [1] dataset. We filter out inputs of length greater than 20000 characters to satisfy our memory constraints. We use the standard train-validation-test split. We compare the original model architecture with Set-LLM with single runs (no adversarial setting). Our initial results do not include hyperparameter tuning, we use the best-performing hyperparameters from the multiple-choice experiments. The experiments will be expanded for the final paper.

We report standard metrics, this includes Rouge F1 scores and compression - the length of the original text divided by the length of the summary. We also include Rouge F1 scores between the target summaries and the model summaries (R1*, R2*, R3*).

Set-LLM outperforms the finetuned baseline with all three base models in all metrics. However, the margin is significantly narrower with the larger Gemma 7B model. Note that the compression rate is always lower for Set-LLM, meaning the summaries it produces are longer.

[1] Fabbri, Alexander R., et al. "Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model." arXiv preprint arXiv:1906.01749 (2019).

Multi-Document Question Answering (Retrieval-Augmented Generation)

Motivated by [2], we evaluate Set-LLM on retrieval-augmented question answering. The authors introduce a task mimicking the retrieval-augmented generation setup. They ask questions with 20 supporting documents, where only one document contains the information required for the answer. They show that commercial LLMs (GPT3.5 Turbo) produce a U-shaped performance graph given the location of the answer. If the document with the answer is listed first or last, then the LLM answers the question correctly with a much higher probability (e.g. 75% in position 1/20 vs 55% in position 10/20). In contrast, Set-LLM is guaranteed to have a flat performance curve, since reordering the documents does not affect the outcome.

We split the dataset into 60% training, 10% validation, and 30% test data. We replicate the training data 5 times and shuffle the documents randomly for each input. This is done to avoid the model overfitting to a specific formulation of each training question. We have not conducted hyperparameter tuning yet, the hyperparameters were taken from the multiple-choice experiments.

We have 5 different test sets with the answer document placed in 5 different locations (first, 5th, 10th, 15th, last). We provide the results for two LLama-based models below. All scores denote the percentage accuracy.

[2] Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2024. Lost in the Middle: How Language Models Use Long Contexts. Transactions of the Association for Computational Linguistics, 12:157–173.

W3: Including baselines that do not provide input choices (i.e., original dataset prompt) seems tangential to motivation and are not sufficiently analyzed in the experimental section.

This is a great suggestion. We have run experiments with a pretrained model and with a model (pre-)finetuned on the ultrafeedback dataset and then finetuned on the multiple choice datasets separately. We finetuned these models again, because the dataset now doesn't include the choices in the question. We report results below for Gemma 2B, but we will include results for all models in the final paper. There is no adversarial setting for this baseline. We include Set-LLM scores for comparison. We see that our model clearly outperforms this additional baseline.