Divide, Reweight, and Conquer: A Logit Arithmetic Approach for In-Context Learning
Abstract: In-Context Learning (ICL) emerges as a key feature for LLMs, allowing them to adapt to new tasks by leveraging task-specific examples without updating model parameters. However, ICL faces challenges with increasing numbers of examples due to performance degradation and quadratic computational costs. In this paper, we propose Logit Arithmetic Reweighting Approach (LARA), a novel framework that enhances ICL by using logit-based ensembling of multiple demonstrations. Our approach divides long input demonstrations into parallelizable shorter inputs to significantly reduce memory requirements, and then effectively aggregate the information by reweighting logits of each group via a non-gradient optimization approach. We further introduce Binary LARA (B-LARA), a variant that constrains weights to binary values to simplify the search space and reduces memory usage by filtering out less informative demonstration groups. Experiments on BBH and MMLU demonstrate that LARA and B-LARA outperform all baseline methods in both accuracy and memory efficiency. We also conduct extensive analysis to show that LARA generalizes well to scenarios of varying numbers of examples from limited to many-shot demonstrations.
- In-context learning with retrieved demonstrations for language models: A survey. ArXiv preprint, abs/2401.11624, 2024. URL https://arxiv.org/abs/2401.11624.
- Language models are few-shot learners. In Hugo Larochelle, Marc’Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https://proceedings.neurips.cc/paper/2020/hash/1457c0d6bfcb4967418bfb8ac142f64a-Abstract.html.
- Controlled text generation via language model arithmetic. ArXiv preprint, abs/2311.14479, 2023. URL https://arxiv.org/abs/2311.14479.
- A survey for in-context learning. ArXiv preprint, abs/2301.00234, 2023. URL https://arxiv.org/abs/2301.00234.
- The llama 3 herd of models. ArXiv, abs/2407.21783, 2024. URL https://api.semanticscholar.org/CorpusID:271571434.
- Semantic compression with large language models. 2023 Tenth International Conference on Social Networks Analysis, Management and Security (SNAMS), pp. 1–8, 2023.
- Connecting large language models with evolutionary algorithms yields powerful prompt optimizers. ArXiv preprint, abs/2309.08532, 2023. URL https://arxiv.org/abs/2309.08532.
- Adapting arbitrary normal mutation distributions in evolution strategies: the covariance matrix adaptation. Proceedings of IEEE International Conference on Evolutionary Computation, 1996.
- Structured prompting: Scaling in-context learning to 1, 000 examples. ArXiv preprint, abs/2212.06713, 2022. URL https://arxiv.org/abs/2212.06713.
- In-context learning creates task vectors. ArXiv preprint, abs/2310.15916, 2023. URL https://arxiv.org/abs/2310.15916.
- Measuring massive multitask language understanding. In Proc. of ICLR. OpenReview.net, 2021. URL https://openreview.net/forum?id=d7KBjmI3GmQ.
- Lorahub: Efficient cross-task generalization via dynamic lora composition. ArXiv preprint, abs/2307.13269, 2023. URL https://arxiv.org/abs/2307.13269.
- Offset unlearning for large language models. ArXiv preprint, abs/2404.11045, 2024. URL https://arxiv.org/abs/2404.11045.
- Mistral 7b. ArXiv preprint, abs/2310.06825, 2023a. URL https://arxiv.org/abs/2310.06825.
- Llmlingua: Compressing prompts for accelerated inference of large language models. In Conference on Empirical Methods in Natural Language Processing, 2023b.
- In-context learning with many demonstration examples. ArXiv preprint, abs/2302.04931, 2023. URL https://arxiv.org/abs/2302.04931.
- Long-context llms struggle with long in-context learning. ArXiv preprint, abs/2404.02060, 2024. URL https://arxiv.org/abs/2404.02060.
- Contrastive decoding: Open-ended text generation as optimization. In Annual Meeting of the Association for Computational Linguistics, 2022.
- Tuning language models by proxy. ArXiv preprint, abs/2401.08565, 2024. URL https://arxiv.org/abs/2401.08565.
- What makes good in-context examples for GPT-3? In Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, pp. 100–114, Dublin, Ireland and Online, 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.deelio-1.10. URL https://aclanthology.org/2022.deelio-1.10.
- Versatile black-box optimization. Proceedings of the 2020 Genetic and Evolutionary Computation Conference, 2020.
- In-context vectors: Making in context learning more effective and controllable through latent space steering. ArXiv preprint, abs/2311.06668, 2023. URL https://arxiv.org/abs/2311.06668.
- Fine-tuning language models with just forward passes. ArXiv preprint, abs/2305.17333, 2023. URL https://arxiv.org/abs/2305.17333.
- Gemma: Open models based on gemini research and technology. ArXiv preprint, abs/2403.08295, 2024. URL https://arxiv.org/abs/2403.08295.
- Learning to compress prompts with gist tokens. ArXiv preprint, abs/2304.08467, 2023. URL https://arxiv.org/abs/2304.08467.
- Llmlingua-2: Data distillation for efficient and faithful task-agnostic prompt compression. ArXiv preprint, abs/2403.12968, 2024. URL https://arxiv.org/abs/2403.12968.
- Ingo Rechenberg. Evolutionsstrategie : Optimierung technischer systeme nach prinzipien der biologischen evolution. 1973. URL https://api.semanticscholar.org/CorpusID:60975248.
- Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. ArXiv preprint, abs/2206.04615, 2022. URL https://arxiv.org/abs/2206.04615.
- Function vectors in large language models. ArXiv preprint, abs/2310.15213, 2023. URL https://arxiv.org/abs/2310.15213.
- Rationale-augmented ensembles in language models. ArXiv preprint, abs/2207.00747, 2022. URL https://arxiv.org/abs/2207.00747.
- Emergent abilities of large language models. ArXiv preprint, abs/2206.07682, 2022. URL https://arxiv.org/abs/2206.07682.
- Prompt compression and contrastive conditioning for controllability and toxicity reduction in language models. In Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 5621–5634, Abu Dhabi, United Arab Emirates, 2022. Association for Computational Linguistics. URL https://aclanthology.org/2022.findings-emnlp.412.
- Effective long-context scaling of foundation models. ArXiv preprint, abs/2309.16039, 2023. URL https://arxiv.org/abs/2309.16039.
- Recomp: Improving retrieval-augmented lms with compression and selective augmentation. ArXiv preprint, abs/2310.04408, 2023a. URL https://arxiv.org/abs/2310.04408.
- Fwdllm: Efficient fedllm using forward gradient. arXiv preprint arXiv:2308.13894, 2023b.
- Safedecoding: Defending against jailbreak attacks via safety-aware decoding. ArXiv preprint, abs/2402.08983, 2024. URL https://arxiv.org/abs/2402.08983.
- Dpzero: Private fine-tuning of language models without backpropagation. In Forty-first International Conference on Machine Learning, 2024.
- Manifold neural network with non-gradient optimization. IEEE Transactions on Pattern Analysis and Machine Intelligence, PP:1–1, 2022.
Paper Prompts
Sign up for free to create and run prompts on this paper using GPT-5.
Top Community Prompts
Collections
Sign up for free to add this paper to one or more collections.