LLM Augmented LLMs: Expanding Capabilities through Composition

TL;DR

CALM tackles the challenge of enhancing large foundational language models by composing frozen anchor LLMs with specialized augmenting models through a compact cross-attention interface. It learns a small projection and cross-attention over selected layers to fuse representations, enabling new capabilities such as low-resource language translation and code understanding without fine tuning the anchor. Across KV arithmetic, low-resource language inclusivity, and code tasks, CALM delivers consistent gains over the baselines and outperforms parameter efficient tuning methods, while avoiding catastrophic forgetting. The approach offers a practical, data-efficient path to extend the capabilities of large models by reusing existing domain experts as augmentors.

Abstract

Foundational models with billions of parameters which have been trained on large corpora of data have demonstrated non-trivial skills in a variety of domains. However, due to their monolithic structure, it is challenging and expensive to augment them or impart new skills. On the other hand, due to their adaptation abilities, several new instances of these models are being trained towards new domains and tasks. In this work, we study the problem of efficient and practical composition of existing foundation models with more specific models to enable newer capabilities. To this end, we propose CALM -- Composition to Augment Language Models -- which introduces cross-attention between models to compose their representations and enable new capabilities. Salient features of CALM are: (i) Scales up LLMs on new tasks by 're-using' existing LLMs along with a few additional parameters and data, (ii) Existing model weights are kept intact, and hence preserves existing capabilities, and (iii) Applies to diverse domains and settings. We illustrate that augmenting PaLM2-S with a smaller model trained on low-resource languages results in an absolute improvement of up to 13\% on tasks like translation into English and arithmetic reasoning for low-resource languages. Similarly, when PaLM2-S is augmented with a code-specific model, we see a relative improvement of 40\% over the base model for code generation and explanation tasks -- on-par with fully fine-tuned counterparts.

Figures (2)

• Figure 1: Overview of CALM. To augment an anchor LLM ($\mathbf{m}$$_{\text{B}}$) with new capabilities through composition with a specialized augmenting model ($\mathbf{m}$$_{\text{A}}$). Figure illustrates three $\mathbf{m}$$_{\text{A}}$ with different capabilities: key-value mapping (left), low-resource languages (center), and code (right). Models $\mathbf{m}$$_{\text{A}}$ and $\mathbf{m}$$_{\text{B}}$ remain unchanged () during composition. A few additional parameters are learnt over models' layer representations. Leftmost plot shows an $\mathbf{m}$$_{\text{A}}$ trained on a set of string-integer mappings, e.g., {$x_1: 10$, $\dots$, $x_n: 2$}. $\mathbf{m}$$_{\text{B}}$ is a large LM with arithmetic capabilities. CALM composes these two frozen models to solve the task of arithmetic on keys which either models could not solve on their own (\ref{['sec:kv_experiments']}). Notably, CALM generalizes to the entire key-value set despite training with arithmetic examples spanning only 20% of the keys.
• Figure 2: Gains seen by the composed model $\mathbf{m}$$_{\text{A}\oplus\text{B}}$ over the anchor model, $\mathbf{m}$$_{\text{B}}$, for the complete set of FLORES-200 languages. The languages are sorted from low to high-resource.