Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

All Language Models Large and Small

Zhixun Chen, Yali Du, David Mguni

TL;DR

All Language Models Large and Small presents LONDI, a framework that learns when to call a costly DEEPTHINK LM by coupling it with a fast QUICK LM through a switching-control RL agent. The core idea is to formulate activation as a Markov decision process and to optimize a policy that balances task performance against computational cost via a per-activation penalty $c$, with a budgeted variant LONDI-B handling explicit DEEPTHINK call limits. The authors prove convergence to the optimal activation policy and demonstrate substantial resource savings (up to 30% GPU reduction) on ScienceWorld and BabyAI-Text while maintaining strong task performance. They also show that LONDI-B can operate as a plug-and-play framework, adaptable to different LMs and encoders, and provide a budget-aware extension that preserves performance within fixed DEEPTHINK call budgets. The work has practical implications for deploying capable LMs on resource-constrained devices by enabling selective outsourcing of expensive reasoning to larger models only where needed.

Abstract

Many leading language models (LMs) use high-intensity computational resources both during training and execution. This poses the challenge of lowering resource costs for deployment and faster execution of decision-making tasks among others. We introduce a novel plug-and-play LM framework named Language Optimising Network Distribution (LONDI) framework. LONDI learns to selectively employ large LMs only where complex decision-making and reasoning are required while using low-resource LMs (i.e. LMs require less GPU usage, but may not be able to solve the problem alone) everywhere else. LONDI consists of a system of two (off-)policy networks, an LM, a large LM (LLM), and a reinforcement learning module that uses switching controls to quickly learn which system states to call the LLM. We then introduce a variant of LONDI that maintains budget constraints on LLM calls and hence its resource usage. Theoretically, we prove LONDI learns the subset of system states to activate the LLM required to solve the task. We then prove that LONDI converges to optimal solutions while also preserving budgetary constraints on LLM calls almost surely enabling it to solve various tasks while significantly lowering computational costs. We test LONDI's performance in a range of tasks in ScienceWorld and BabyAI-Text and demonstrate that LONDI can solve tasks only solvable by resource-intensive LLMs while reducing GPU usage by up to 30%.

All Language Models Large and Small

TL;DR

All Language Models Large and Small presents LONDI, a framework that learns when to call a costly DEEPTHINK LM by coupling it with a fast QUICK LM through a switching-control RL agent. The core idea is to formulate activation as a Markov decision process and to optimize a policy that balances task performance against computational cost via a per-activation penalty , with a budgeted variant LONDI-B handling explicit DEEPTHINK call limits. The authors prove convergence to the optimal activation policy and demonstrate substantial resource savings (up to 30% GPU reduction) on ScienceWorld and BabyAI-Text while maintaining strong task performance. They also show that LONDI-B can operate as a plug-and-play framework, adaptable to different LMs and encoders, and provide a budget-aware extension that preserves performance within fixed DEEPTHINK call budgets. The work has practical implications for deploying capable LMs on resource-constrained devices by enabling selective outsourcing of expensive reasoning to larger models only where needed.

Abstract

Many leading language models (LMs) use high-intensity computational resources both during training and execution. This poses the challenge of lowering resource costs for deployment and faster execution of decision-making tasks among others. We introduce a novel plug-and-play LM framework named Language Optimising Network Distribution (LONDI) framework. LONDI learns to selectively employ large LMs only where complex decision-making and reasoning are required while using low-resource LMs (i.e. LMs require less GPU usage, but may not be able to solve the problem alone) everywhere else. LONDI consists of a system of two (off-)policy networks, an LM, a large LM (LLM), and a reinforcement learning module that uses switching controls to quickly learn which system states to call the LLM. We then introduce a variant of LONDI that maintains budget constraints on LLM calls and hence its resource usage. Theoretically, we prove LONDI learns the subset of system states to activate the LLM required to solve the task. We then prove that LONDI converges to optimal solutions while also preserving budgetary constraints on LLM calls almost surely enabling it to solve various tasks while significantly lowering computational costs. We test LONDI's performance in a range of tasks in ScienceWorld and BabyAI-Text and demonstrate that LONDI can solve tasks only solvable by resource-intensive LLMs while reducing GPU usage by up to 30%.
Paper Structure (26 sections, 7 theorems, 24 equations, 7 figures, 7 tables, 2 algorithms)

This paper contains 26 sections, 7 theorems, 24 equations, 7 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. LONDI
  4. Markov Decision Process
  5. Problem Formulation
  6. Switching Controls
  7. Architecture
  8. Convergence and Optimality of LONDI
  9. LONDI with a DEEPTHINK Call Budget
  10. Algorithm
  11. Experiments
  12. Results and Analysis
  13. Conclusion
  14. Schematic of LONDI
  15. Ablation Study
  16. ...and 11 more sections

Key Result

Theorem 1

i) Let $v_S:\mathcal{S}\to\mathbb{R}$ then for any fixed policies ${\pi}^{\rm QUICK},{\pi}\in {\Pi}$ the solution of Switcher's problem is given by $\underset{k\to\infty}{\lim}T_S^kv_S(\cdot|{\pi},\mathfrak{g})=\underset{\hat{\mathfrak{g}}}{\max}\;v_S(\cdot|{\pi},\hat{\mathfrak{g}})$, where $T_S$ is

Figures (7)

  • Figure 1: (a). The ScienceWorld task, Create a circuit. To complete the task, the agent must navigate to the hallway first and then determine the correct room to enter then use the material present in the room. (b),(c),(d),(e) sub-figures are the heatmap of LONDI DEEPTHINK calls with different budget on ScienceWorld task, Create a circuit. Since the agent needs to pass through the hallway to reach the workshop initially, LONDI must activate DEEPTHINK model in the hallway at least once. Subsequently, the remaining budget allocated to activating the DEEPTHINK model at the workshop to accomplish the task and obtain a higher reward. Therefore, as the budget decrease, LONDI focuses its activations of DEEPTHINK solely at the critical parts of the map such as hallway, resulting in lighter colors in the heatmap.
  • Figure 2: An illustration of one BabyAI task, PutNextLocal:"put the blue key next to the green ball". The shadow area represents the observable space of the agent.
  • Figure 3: Performance of LONDI on ScienceWorld task: Identify Longest-then-shortest-lived animal with different cost (normalized)
  • Figure 4: Performance of LONDI on BabyAI-Text with mixed tasks. The cost and reward are normalised values.
  • Figure 5: The schematic of LONDI. The Diamond represents the decision point, the square means a process or action, and the oval-like shape means data storage. The Switcher agent receives an environmental observation and makes a decision on which LLM module to utilise based on factors such as switch state, switch probability, and observation. The transition is then stored in the buffer for training the Switcher policy in subsequent iterations.
  • ...and 2 more figures

Theorems & Definitions (13)

  • Theorem 1
  • Proposition 1
  • Theorem 2
  • Definition 1
  • Definition 2
  • Lemma 1
  • Lemma 2
  • proof
  • Lemma 3
  • Theorem 3: Theorem 1, pg 4 in jaakkola1994convergence
  • ...and 3 more