LLM Augmented LLMs: Expanding Capabilities through Composition
Authors: Rachit Bansal
Bidisha Samanta
Siddharth Dalmia
Nitish Gupta
Sriram Ganapathy
Abhishek Bapna
Prateek Jain
Partha Talukdar
Venue: ICLR 2024 poster (2024)
Keywords: Large Language Models
Model Composition
Knowledge Augmentation
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Foundational models with billions of parameters which have been trained on large corpus 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.

LLM AUGMENTED LLMs: EXPANDING CAPABILITIES THROUGH COMPOSITION
Rachit Bansal' Bidisha Samanta' Siddharth Dalmia² Nitish Gupta' Shikhar Vashishth Sriram Ganapathy' Abhishek Bapna' Prateek Jain' Partha Talukdar' !Google Research India ?Google DeepMind
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. How- ever, due to their monolithic structure, it is challenging and expensive 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 capa- bilities. To this end, we propose CALM- CLLM-ooopsittion to Augment Language Models-which introduces crossaattention between models to compose their rep- resentations and enable new capabilities. Salient features of CALM are: (i) Scales up LLMS on new tasks by 're-using' existing LLMs along with few additional parameters and data, (ii) Existing model weights are kept intact, and hence pre- serves 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 trans- lation into English and arithmetic reasoning for low-resource languages. Simi- larly, when Pal M2-S is augmented with a code-specific model, we see a relative improvement of 40% over the base model for code generation and explanation with fully fine-tuned counterparts.
INTRODUCTION
Large Language Models (LLMs) have shown to encompass a range of foundational capabilities such as commonsense and factual reasoning. world knowledge, and coherent language generation (Bubeck et al., 2023; Google et al., 2023). Leveraging these foundational capabilities, number of efforts in the community have fine-tuned these models to enable domain-specific capabilities such as code generation, copy editing, and mathematical problem solving (Lewkowyez et al., 2022; Singhal et al., 2023). This has resulted in the development of several specialized large models with domain- specific capabilities. For example, there models that do well on standard code generation but are not as proficient in general logical reasoning and vice-versa. Presence of such large number of domain-specific models leads to a natural question: Can we compose an anchor model with a domain-speciff augmenting model to enable new capabilities? For example, can we compose an augmenting model's code understanding capability with an anchor LLM's language generation capability to enable code-to-text generation capability?
The typical approach for this problem to further pre-train or (efficiently) fine-tune the anchor model on the data that was originally used to train the augmenting model (Hu et al., 2021; Kessler et al., 2022). However, many a times such solutions are not feasible since training large models is computationally expensive, especially since the augmenting model itself may be an LLM trained on massive corpora. Further, processing data from multiple sources might not be feasible due to privacy concerns and organizational boundaries. Working with multiple distinct models is also desirable since allows the reuse of existing models with established capabilities, providing better control and avoiding catastrophic forgetting that is prevalent in conventional approaches.
Correspondence Rachit and Bidisha: [brachit, bidishasamanta| @ google.com

---

Since x==10, xs-14, x,m2, X1 +:x 38
Keyvalue Mapping
Numeric Artthmetic
Everything but the kitchen sink
Implements the classic word game of Hangman
180++
urce Language Pre-traite
*
Pre-trained GitHuub
+ me
m4
+
Attention
mi
WK
WN
Wq
M
lui
TRI
from En: Source Sentences
What this Python code do? ython Code Srippet>
What is the value of +x x.?
Figure 1: Overview of CALM. To augment anchor LLM (mg) with new capabilities through composition with a specialized augmenting model (mą)) Figure illustrates three mA with differ- ent capabilities: key-value mapping (leff), low-resource languages (center), and code (right). Mod- els mA and mg remain unchanged (2E) during composition. A few additional parameters are learnt over models' layer representations. Leftmost plot shows an mA trained on a set of string-integer mappings, c.g., {x1 10, In 2}. mg is a large LM with arithmetic capabilities. CALM com- poses these two frozen models to solve the task of arithmetic on keys which either models could not solve on their own (84.1). Notably, CALM generalizes to the entire key-value set despite training with arithmetic examples spanning only 20% of the keys.
To address the training and the data challenges mentioned above, we propose and study a practical setting for model composition: (i) we are given access one (or more) augmenting model(s) and an anchor model, (ii) we are not allowed to modify the weights of either models, and (iii) we only have access to small amount of data, representing the "combined skills" of the given models, e.g., code generation with complex logical reasoning.
Prior work has largely approached the question of composition from either a routing or a merging standpoint, neither of which provide effective solution capture this setting. Routing between the given models, i.e., choosing an output of one model over the other (Ma et al., 2019), or performing a soft ensemble (Muqeeth et al., 2023) is not effective when neither of the models can demonstrate the desired capability. Another body of work creates a combined model by an arithmetic combination of base model parameters (Wortsman et al., 2022; Ilharco et al., 2022; Matena & Raffel, 2022). However, these settings are naturally restrictive and their efficacy is unclear when combining models with different sizes and pre- pre-training objectives (Yadav et al., 2023).
In this work, we propose a novel Composition to Augment Language Models (CALM) framework to address the general model composition setting mentioned above. Rather than shallow combi- nation of the augmenting and anchor LMs (Wortsman et al., 2022; Ilharco et al., 2022), CALM introduces small number of trainable parameters over both augmenting and anchor models' inter- mediate layer representations. CALM finds an effective combination of the given models to perform new challenging tasks more accurately than either of the models alone, while preserving the capa- bilities of individual models. Figure highlights few motivating scenarios for CALM.
We study key practical applications of CALM: language inclusivity and code generation. For lan- guage inclusivity ($4.2), we use model that has been trained on a set of low-resource languages. We observe that composing this model with the L.LM allows us borrow its generation and reason- ing capabilities to achieve significantly better performance on translation and arithmetic reasoning tasks for low-resource languages (Tables 2 and 3). This composed model outperforms not only the two base models but also versions of the LLM that have been further pre-trained or LoRA (Hu al., 2021) finettuned for the set of low-resource languages. For code generation (G4.3), we use a model open-source code across variety of programming languages. Compos- ing this model with the LLM-henne borrowing its low-level logic and generation capabilities- outperforms the two base models (Table 4) on code explanation and code completion tasks.
that has been trained
2

---

2 RELATED WORKS
Parameter efficient fine-tuning: A large body of prior work focuses on parameter efficient ways of fine-tuning models for new tasks by introducing a small number of trainable parameters, keeping the original model intact (Houlsby et al., 2019; Wang et al., 2020; Pfeiffer et al., 2021; Hu et al., 2021: Kessler et al., 2022). Since this paradigm allows a small set of new parameters be trained, it is challenging to use these approaches to augment novel domains and knowledge sources that are entirely absent from the original training corpus. contrast, CALM enables a model be adapted to new domains using augmenting models. In Section 4.4, we draw empirical comparisons between CALM and LoRA (Hu et al., 2021), ) representative parameter efficient fine-tuning method.
Model Merging: Merging different expert models with simple techniques like task vector aver- aging provides a way of recombining different capabilities of these models (Ilharco et al., 2022; Matena & Raffel, 2022). However, these methods are only relevant when the original models well aligned. Other related approaches are also applicable only when the models are derived from the same model (Matena & Raffel, 2022) they are of same size (Mugeeth al., 2023). In contrast, CALM is more generic and is applicable to any set of models.
Model and Task Compositionality: The modular encoder-decoder based method in Dalmia et al. (2022) adapts components of encoder-decoder models to allow flexible re-usability of different en- coders, each with their own capabilities. Several past studies explore compositionality of modality- specific encoders with language models to serve multi-modal use-cases (Ziegler et al., 2019; Alayrac et al., 2022). Typically, they introduce cross-attention parameters across a language model order to attend representations from image encoder and show an effective transfer of modalities across models. In this work, we extend the ideology of model re-use and modularity to composition of capabilities in large language models.
Models as Tools: Another interesting direction for using multiple language models to solve downstream task has been to perform composition in the models' input text space (Zeng et al., 2022; Shen et al., 2023). Schick et al. (2023) have demonstrated how a model can be taught to use external toele-there might be an opportunity to investigate if other models can be called as a part of the same framework. Since these approaches require a large amount of prompt engineering, in this work we focus on composition through representations that can be learnt automatically.
3 COMPOSITION TO AUGMENT LANGUAGE MODELS (CALM)
Given an anchor model my and an augmenting model m, CALM aims to compose the two models (mA@n) to enable new capabilities composition of capabilities of the two individual models.
As discussed in the introduction, we study this composition in a practical setting with the following assumptions: i) we can access weights, run forward and backward pass, and access intermediate representations of both ma and mA ii) we are not allowed to change weights both the models, iii) we do not have access the training data, hyperparameters, training states of both the base models, iv) we are provided a few examples from the target composition domain.
The goal is to learn composition mA@B = f(m,, mg, Oc. D&)) to achieve some joint task C. The weights of ma and mg are frozen. θc is the additional set of trainable parameters introduced to learn the composition and D. refers to the examples that are used to learn this composition.
3.1 LEARNING TO COMPOSE (ec)
As outlined in Figure 1, we operate over selected set of layers from mg and mA at all times. We learn two sets of additional parameters over these layers: (i) A simple set of linear transformations, fpag(.) that maps hh layer representation from m, the dimensionality of representations from my. and (ii) A set of cross-attention layers, fa.sss...) that cross-attend between this transformed layer representation and a 7th layer representation from mB.
Compositional Layers: Let the augmenting model mA and the anchor model ma have NA and NB layers, respectively. Also, let DA and D3 be the token dimensionality of the models. We
3

---

first choose set of compositional layers- LA and Lg for both models, over which the set of new learnable parameters are introduced during composition. n, = A and ng = [LB] For simplicity, set RA n n and the gap between two contiguous selected layers is kept uniform based on the number of selected layers-that is, (22 - 4) = ... (Ln ((--1) = N/n. Further, A (HAI, Haz...., HAn, denote the layer representation of a given imput after each layer in LA.
Learned Projections: Next map representations from mA to that of mp via projection layer. In particular, for each layer in LA, we learn a projection function fproj IR + R"n that projects representations from these layers to the desired representation size of mg. Let,
Sproj AI & {Sproj (HAI), Sprij (HA2)...... Sproj (HAn)))
This transformation enables cross-attention across models, and also performs an alignment of rep- resentations from mA and mg despite frozen weights of the base models.
Cross-attention Layers: Similar to the multi-headed cross-attention in encoder-decoder models (for example Vaswani et al. (2017) and Raffel al. (2020)) we introduce aross-attention between representations of the anchor and the augmenting model. In particular, we use fpooj (HAi) from the augmenting model as the key and value vectors each head in cross-attention. We use the vector HB, from the anchor model as the query vector, which leads to the following cross-attention setup:
feross (Frooj (HAr), HBj) Concat, (heads) Wo Ak Nu where, head,, = Attn. (Qn, K. Va), and, Q HRj W KA V = /pejj (HAr) W (HA:)
Here, Nи represents the number of attention heads used for cross-altention which, in our case, is typically the same as the number of heads used for self-attention in mg. Each of W E R DnxDn, are learnable weight matrices, where k E (1..NIr).
wº. WK E and W. DDDxxDa//ie
Finally, the cross-attention output is added as residual connection the layer representations of mg. The resultant output vector, in-turn, is the input to the succeeding layer in mp:
HASB = HBy + feross (froo H HA.). HB;)
Here, HA@By denotes the input to the (j + 1) 1)th layer of the composed model. All layers in L and LB are utilized in similar manner. Propagating over the remaining layers in my gives us a final output token Y/ decoded for the tt timestep. Akin usual auto-regressive decoding, the output token for each time-step is appended to the input: Iz+1 Ia @ yt, Since the updated input at each time step is passed both models, all representations for the two models are refreshed.
3.2 COMPOSITION TRAINING DATA (Dc)
Since the target model mA@B involves composition over the two models mA and mg, we construct the set of training examples De to depict a "combined skill" that enables Oc to attend over the two models appropriately for the target task.
Ideally, if the set of tasks involved in composition task are distinguished as t1 and t, respectively, then we design Dc depict a joint task C. For example, with respect to our synthetic key-value setup: our final task (C) is to perform arithmetic over a set of keys. The augmenting model m, is trained learn the given key-value pairs (notated as task, t1) and the anchor model mp is generic model that can perform numeric arithmetic well (task t²). For learning the set of parameters Oc for composition, we consider Dg to be arithmetic over a held-in set of keys (task C), encompassing combined skills from the models. In contrast fine-tuning approaches like LoRA (Hu et al., 2021) that would require the entire knowledge source (here, key-values) during training time, we find that training composition on only fraction of the keys can generalize the full set.
In other real world settings, a clear distinction in specializing tasks for each model might be difficult to formulate and hence defining task that captures the combined skills can be challenging. We find

---

that using a of examples that capture certain capabilities of the two models suffices, i..., some rough notion of tAuB- For our language inclusivity task, we use a mixture of examples containing small amount of low-resource language and high-resource language data.
Finally, we note that while the method has been presented for a setting with one anchor model and only one augmenting model, CALM is applicable multiple augmenting models as well. In particular, CALM would require learning similar projection and cross-attention components between the anchor and each of the augmenting model. We leave a thorough investigation of this as topic of future work.
Composing multiple models:
4 EXPERIMENTS
We demonstrate the following in three domains: (a) an anchor LLM (mg) can be composed with an augmenting model (mA) trained on mappings between string keys and number values to solve arith- metic expressions over those keys requiring both, knowledge of the KV mappings and arithmetic capabilities ($4.1), ;; (b) how CALM can be used expand the language coverage of an anchor LLM (m₂) to low-resource languages it has not seen during pre-training. We show that an augmenting model (mx) pre-trained on lowrresource languages can be composed with such an anchor model to significantly improve translation and math-word problem solving capabilities in low-resource lan- guages ($4.2); (c) how code completion and explanation can be improved by composing an anchor LLM with an augmenting model (MA₂ specializing the code domain (54.3).
In all experiments, we start with a PaLM2-XXS model and further train it on fomain-specific data to arrive at augmenting model (ma)) that is then kept frozen during composition. Note that task specific training data was used to train CALM. We use PaLM2-XS or PaLM2-S models as the anchor LLM (mg) that is also kept frozen during composition training. For all our experiments, set NA/n = 4, i.e., we perform composition using every 4th layer output from ma. Correspondingly. layers from mp (Ln) are chosen such that n T n, hence n Nn/4.
4.1 KEY-VALUE ARITUMETIC
We first study the setting where we have a small augmenting LM that has been trained to memorize turing-to-integer key-value (KV) mappings, and a large anchor LM that is capable of performing arithmetic over integers. We wish to use CALM to compose them and enable a new capability of solving arithmetic expressions containing those keys.
Key- Value Domain Knowledge We first generate repository of KV pairs containing NKv 25K pairs by sampling English strings of 6 characters from the vocabulary of the PaLM2-XXS model and randomly assigning them unique integer values in the range [1, NKv]. This constitutes the knowl- edge artifact, DRv. We further generate a collection of arithmetic expressions (DKv.exp) containing addition (+), subtraction (-), and multiplication
6 keys by randomly sampling keys from DRv and operations to perform between them. We generate three datasets:
(x) operations between
(i) K KV-Substitution (DRv:suns): This dataset maps each expression in DRV to an expression where the keys are replaced by their corresponding values. For example, this dataset contains exam- ples of the form (<K1>< <K2> <K3>, 10 22 - 24)..
(ii) KV-Arithmetic (DRvmath): This dataset maps each expression in DRVEXP the numeric value arrived at by solving the arithmetic expression when the keys would be replaced by the correspond- ing values. For example, examples in this dataset look like (<K1> <K2> <K3>, 8).
(DxUM-MATH): This dataset maps the value substituted version of each expression in DRveXP to the numeric value arrived at by solving the arithmetic expression. For example, examples in this dataset look like (10 + 22 24, 8).
(iii)
Models We obtain augmenting model mA by further training pre-trained PaLM2-XXS model on
Dgv-SuBS to make memorize the KV pairs in DKv. Note that, training on DKvSSUS does not teach pre-trained PaLM2-XS
this augmenting model how to solve arithmetic expressions. Next, model as the anchor model mg. This model is capable of solving numeric expressions with decent performance (see Table 1). Note that, this model has knowledge of the KV pairs in Dgy.
5

---

We now take examples from the KV-Substitution dataset D DEV-SUBS that only span 20% of the keys in DRv to form the training data for composition (De). We use Dc to compose the augmenting model (ma) having knowledge of Dxy and the pre-trained anchor model mp by training the composition parameters (9c) using CALM as explained in g3. Both mA and mB are kept unchanged.
Evaluation Task We evaluate the composed model maen for its ability to solve arithmetic ex- pressions containing keys from Dgv. Specifically, we evaluate on the subset of DRv.MATT dataset that does not contain expressions used in Dc during training. This way, we are able to measure the composed model's ability to generalize to keys beyond what was observed during training.
Results Table shows the performance of the three models: mA. mg. and MяşB across the aforemen- tioned datasets. First, we observe that the augmenting model mA achieves 98. 1% at the KV-Substitution task showing that memorizes DRy well. Next, we see that performs poorly (4.2%) at the Numeric-Arithmetie task showing that it does not have arithmetic capabili- ties. As result, this model is not able to solve arith- metic expressions containing keys from DRv.
CALM
mA
mg
(MAgB)
D DKV-SUBS
98.1
0.0
92.9
NUMM
4.2
73.7
72.0
DRVMMAT
0.7
0.0
84.3
As expected, the anchor model mp gets 0% accuracy on the KV-Substitution and KV-Arithmetic tasks it has not seen any data from Dxv. However, performs well (73.7%) on the Humeric-Arithmeee task demon- strating capability of arithmetic over numerals.
Table 1: Evaluation (accuracy (%)) for a synthetic key-value (KV) task. mA is trained memorize the KV mappings while mB excels at arithmetic We see that composition mAą8 is able perform arithmetic over held-out keys.
Lastly, we that the composed model mA@B is able to solve all tasks with high accuracy, especially the KV- Arithmetic task (84.3%) for which both the underlying models fail. This shows that the composed model is able to leverage the relevant capabilities from both the augmenting and anchor model to solve a complex task.
4.2 LOW-RESOURCE LANGUAGE INCLUSIVITY
FLORES-200 (XX to En; chrFI)
Model
lij
mr
taq
nn
su
ban
pl
th
min
acm
avg.
PaLM2-XXS
24.0
16.5
21.6
33.3
20.6
2.1
5.3
63.2
44.0
59.8
29.0
NTL (ma)
32.0
21.6
46.9
50.0
40.6
4.1
4.0
63.8
47.8
61.1
37.2
PaLM2-S (mg)
32.6
24.2
44.6
50.8
50.9
5.4
9.5
69.0
61.0
68.6
41.7
CALM (MA@B)
44.1
30.4
55.1
54.6
54.4
11.8
11.3
69.4
61.1
68.9
46.1
mp | NNL (m)T))
48.1
39.1
59.2
57.5
57.3
11.4
9.9
69.4
61.4
69.0
48.2
Table 2: Translation performance for XX to English direction on the FLORES-200 dataset (Costa- jussà et al., 2022): We show results for a subset of 10 low-resource languages. Note that the com- posed model mAAB significantly outperforms both mA and mg. On the complete language list, mAAB outperforms both the underlying models for 175 of 192 languages (Appendix A: Figure 2). NTL represents skyline where mg has been further pre-trained on DNTL- The composed model ma achieves similar performance for tiny fraction of the training cost.
In this section, we study we can compose such large anchor LM mg with smaller augmenting LM that has been pre-trained on low-resource languages, to perform translation and math-word problem solving tasks presented in these low-resource languages.
Low-resource Language Corpora
We use the long-tail language set and the associated corpora from the Next Thousand Languages (NTL.) effort (Caswell al., 2020; Bapna al., 2022) as the

---

GSM8K (Low-resource Languages; Accuracy)
Model
meo
mfa
pem
efi
min
ilo
ady
mai
nsO
mzn
avg.
PaLM2-XXX
5.2
6,8
6.8
4.0
5.6
7.2
6.0
3.6
7.2
6.8
5.9
+ NTL (m^))
7.6
4.0
4.4
3.2
6.0
4.8
6.4
3.2
6.0
4.8
5.0
PaLM2-S (mg)
28.8
14.0
34.4
14.8
25.2
14.8
30.0
22.8
8.4
31.6
22.5
CALM (mAgg)
34.0
17.6
33.6
18.0
23.6
16.8
36.4
24.8
8.4
36.4
25.0
m
33.2
20.4
31.6
14.0
24.8
14.0
29.2
21.2
9.6
27.6
22.6
(High-resource Languages)
Model
en
tc
bn
SW
ja
zh
th
fr
cs
de
avg.
PaLM2-XXS
5.6
4.0
2.0
7.6
2.0
4.4
6.0
6.8
5.6
9.2
5.3
+ NTL (ma)
4.8
3.6
3.2
4.8
3.2
7.6
6.4
9.2
5.6
7.2
5.6
PaLM2-S (ma)
36.8
19.2
23.2
16.0
2.0
39.2
29.6
38.0
32.4
43.2
28.0
CALM (mA@8)
37.2
28.0
27.2
18.0
2.4
43.6
33.2
42.8
36.0
49.2
31.8
m
NTL
36.0
17.6
18.4
14.4
0.8
33.6
27.2
34.8
31.2
42.0
25.6
Table 3: Evaluations for grade-schou mathematics (GSM) problems on low-resource (LRL) and high-resource (HRL) languages. We observe that CALM yields significant gains for both evaluation sets. Gains on the HRL set suggests that CALM avoids catastrophic forgetting.
domain data DNIL- This large-scale corpora contains web-crawled monolingual sentences and trans- lation pairs for 1000 languages. The dataset has been used for language expansion in translation systems and language models (Garcia et al., 2021; Siddhant et al., 2022).
Models Akin to $4.1, we obtain augmenting model mA by training the PaLM2-XXS model on DNn to impart knowledge about these low-resource languages to the model. For mg. use the pre-trained PaLM2-S model. We use ~ 5% of the same low-resource language corpora DNIL as the training data De compose mA and my via CALM. Since both models are untrained during composition, the anchor model ma is not trained on any of the low-resource language data.
Evaluation Tasks
We evaluate the composed model mA@B on two tasks:
(i) Translating text from a non Ennliinlishh language to English: We carry out these evaluations in 5-shot in-context learning paradigm on the FLORES-200 (Costa-jussà et al., 2022) dataset. This dataset contains examples for 200 high- and low-resource languages.
(ii) Performing grade school math word problems expressed in non-English language: We evaluate on the multilingual version of the GSM-8K dataset (Shi et al., 2023) containing math word problems for English and other high-resource languages. We further generated a silver-standard GSM-8K dataset for low-resource languages by automatically translating the English examples in GSM-8K to 25 low-resource languages supported by Google Translate."
Results Table shows results the FLOREE--20 dataset (Costa-jussà et al., 2022), where the input is low-resource (XX) language sentence and the output should be the corresponding English translation. For 10 low-resource languages shown the Table, we see that both the underlying models mA and mg are outperformed by our composed model mA@B- We find that the composed model mAgs outperforms mg on 175 of the complete set of 192 languages (Appendix A).
Table 3 shows the performance of these models on the grade-school math word problems from the GSM8K task (Cobbe al., 2021) low-resource languages (top) and high-resource languages (Shi et al. (2023); bottom). Firstly, we observe that the augmenting model mA does not perform well on this task due its limited mathematical reasoning capabilities. On the other hand, the anchor model mg does much better given its mathematical reasoning capabilities and transfer-learning from high- resource languages. Finally, we observe that mA@8 outperforms both mA and mB on of 25 low-resource and 9 of 10 high-resource languages, demonstrating effective composition of models.
We perform quality evaluations in Appendix 7.

---

Model
CC (P@1)
T2C (P@1)
C2T (chrF1)
HvmanEval
MBPP
Python
PHP
Go
Java
Ruby
M2
19.5
28.0
28.0
34.7
32.6
29.6
26.5
26.0
Code (MA)
PaLM2-S (mg)
16.4
28.6
30.4
35.5
40.4
31.0
28.8
27.9
CALM (mA@B)
22.5
32.2
30.5
35.8
40.6
31.4
29.3
29.0
m
Code
24.3
43.0
18.9
35.0
41.1
31.1
20.2
27.6
Table 4: Evaluations for code generation and understanding across three tasks: Code Completion (CC), Text-to-Code (T2C), and Code-to-Text (C2T). Augmenting code understanding to mg using mA significantly improves performances across all datasets. mg mf represents a skyline where mg further pretrained on the Doule which shows catastrophic forgetting of text generation task.
See Table 6 (Appendix A.2) for complete of evaluations. Note that the last row in Table 3 shows that my when fine tuned on DNLL leads to worse performance than the pre- pre-trained my indicating forgetting. Composing domain-specific model mA with mp using CALM avoids this.
4.3 CODE UNDERSTANDING AND GENERATION
Code understanding and generation require two distinct types capabilities: (a) knowledge of the syntax and semantics of code, and (b) knowledge of the world that the code is manipulating, While LLMs have wealth of world knowledge, they could often lack the specific knowledge of code syntax due to skewed representation of code data in their pretraining corpora. Conversely, small models trained specifically on code data could exhibit good understanding code syntax, but they may lack broad world knowledge and reasoning. CALM can enable best of both worlds.
Code Domain Data Here, we use the code-specific corpus, Dcade: consisting of open-source code extracted from GitHub heads for variety of programming languages to train m,.
Models Similar to $4.1, version of the PaLM2-XXS model has been further pre-trained on D used mA while the base pre-trained PaLM2-S model acts as mg. We build mx by training CALM with only 7% of the same code data (data used for m,) to have data parity.
Evaluation Tasks
We evaluate the efficacy of CALM on three different tasks:
(i) Todde-Complttinn (CC): Given an initial set of lines of a code, the model prompted to complete the code snippet. Here the aim evaluate the model for code syntax. We perform zero-shot eval- uations on HumanEval benchmark dataset (Chen et al., 202 and report the Pass (P@) I) metric.
(ii) Text-to-Code Co-e (T2C): Given textual context, the model is prompted to generate the correspond- ing code snippet. Here, the evaluation indicates language understanding and code generation capa- bilities. We perform 3-shot inference the MBPP dataset (Austin et al., 2021) and report P@) 1.
(iii) Code-to-Ters (C2T): Given code snippet, the goal is to generate natural language explanation of the code. This task evaluates code understanding and text generation. We perform 3-shot evalua- tions on the CodeXGlue benchmark (Lu et al., 2021) and report chrF1 scores across languages.
Results Table 4 reports comparative performance for the individual models mA and mB, the com- posed version mA@B, and fine-tuned anchor baseline m Code Firstly, evaluations on the HumanEval dataset suggest that mA has superior understanding of code syntax as result of its additional train- ing on Dcode- While, due to the larger scale and general purpose pre-trac pre-training of mg. it excels at general language understanding and hence performs better on the T2C and C2T tasks.
When employing CALM to compose the two models, we observe clear transfer and composition of capabilities through significant performance improvements: 6.1% and 3.6% absolute gains mg on the CC and T2C tasks, respectively. We observe that fine-tuning my on D leads to a significant decline in the C2T performance due catastrophic forgetting. CALM retains the perfor-
8

---

FLORES-200
GSM-8K (LRL)
GSM-8K (HRL)
CC
T2C
C2T
(XX-En)
chrF1 1
#(>m²)
acc.
#(>mB)
#(>mn)
P@
P@ 1
chrF1
A
60.5
175
21.4
20
33.1
11
22.5
32.2
32.6
Vanilla mA
59.2
115
19.0
29.7
8
20.0
28.0
32.2
Random mA
58.8
43
17.8
9
28.5
4
20.1
27.0
32.1
mA as an encoder
59.3
102
19.1
12
29.1
6
16.0
27.0
32.0
LoRA
59.2
82
20.9
15
31.2
9
18.3
28.7
32.6
Table 5: Comparative performance of CALM (mx@n) across various possible ablations. The metric "#(>ma) mg)" depicts the number of languages for which the corresponding model is better than the base for NTL,, mB-out of 192, 25, and 11 languages for the three tasks respectively.
mance and marginally better than mg across all languages. We also study qualitative examples on the C2T task and observe interesting common patterns that are discussed in Appendix B.
4.4 ABLATIONS
Influence of mA We first study the influence of mA by replacing with vanilla and random variants during composition. Table shows the variation of performance across NTL and Code tasks when the specialized mA is replaced with a vanilla PaLM2-XXX checkpoint or an untrained version of the model, i.e., random model. We see that there is a considerable drop of performance with these variants across all tasks. On FLORES-200 XX- En task, languages improved with composition drop to 11 115 and 43 with vanilla and random, respectively. A slight improvement of the vanilla model over ma indicates that an un-specialized model (with different training regime than ma) might have orthogonal capabilities leading to enhanced model. This finding validates that performance gains seen with CALM is result of utilizing mA and not just the additional Oc parameters.
Influence of iterative decoding We also investigate variation where ie., an output token decoded at a given timestep is not amended to m, 's input. In this case, only the prefix representations of mA used. This setting eludes past work for image and text models (Alayrac et al., 2022) where encoder and decoder models are composed. We observe a significant decline in performance across our various tasks when employing this setting.
mA encoder,
Comparision with LoRA Finally, we evaluate a parameter efficient fine-tuning approach by train- ing LoRA (Hu et al., 2021) layers adapt mg. For all experiments, set the LoRA rank such that the number of added parameters is equal to the number of parameters introduced with CALM. We also train LoRA on the same data CALM, i.e., D We see considerable difference in performance between the two approaches across all tasks and metrics.
5 CONCLUSION
The proposed CALM framework composes anchor LLM with specialized augmenting models enable new tasks not achievable by either models individually. CALM does not require updating the individual models and learns a dense interaction between the models through few trainable cross- attention parameters. Our experiments present consistent evidence that CALM learns to utilize the expertise from the two models. That is, when composed with relevant augmenting models, we observe significant uptick in the anchor model's performance across multiple challenging tasks, such as low-resource translation, reasoning, and code explamation/genneaaioo.
That is, CALM is especially useful in scenarios where proprietary data and knowledge is stored in parametric models. With CALM, a foundational LLM could be augmented with such proprietary models to extend variety of foundational capabilities such as reasoning, world knowledge, and coherent generation over the target proprietary domains. Finally, extensions of CALM could be used acquire distinct knowledge from multiple augmenting models.

---

ACKNOWLEDGMENTS
This work was done during RB's pre-doctoral tenure at Google Research, India (GRI) with PT and PJ. RB indebted to Manish Gupta, Divvy Thakkar, and all others who enabled this oppurtunity. RB would also like to thank the members of the Languages team and other researchers at GRI (and beyond), including the incredible pre-doctoral cohort. This work wouldn't have been possible without their constant support. Namely: Aishwarya P.S., Laurent El Shafey, and Qiao Zhang for their massive help in coding and debugging; Palak Jain and Sagar Gubbi for their feedback and support throughout the project; Kartikeya Badola, Shreyas Havaldar, Amandeep Kaur, and Rishabh Tiwari for being the first cars to all ideas; Cyrus Rashtchian and Richa Dixit for their mentorship.
REFERENCES
lean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, lain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katie Millican, Malcolm Reynolds, Roman Ring, Eliza Rutherford, Serkan Cabi, Tengda Han, Zhitao Gong, Sina Samangooei, Marianne Monteiro, Jacob Menick, Sebastian Borgeaud, Andrew Brock, Aida Nematzadeh, Sahand Sharifzadeh, Mikolaj Binkowski, Ricardo Barreira, Oriol Vinyals, Andrew Zisserman, and Karen Simonyan. Flamingo: Visual Language Model for Few-Shot Learning, April 2022. URL http: arxiv org/ abs 14198. arXiv:2204. 14198 [es].
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Temry Quoc Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2 108.07738, 2021.
Ankur Bapna, Isaac Caswell, Julia Kreutzer, Orhan Firat, Daan van Esch, Aditya Siddhant, Meng- meng Niu, Pallavi Baljekar, Xavier Gareia, Wolfgang Macherey, Theresa Breiner, Vera Axelrod, Jason Riesa, Yuan Cao, Mía Xu Chen, Klaus Macherey, Maxim Krikun, Pidong Wang. Alexan- der Gutkin, Apurva Shah, Yanping Huang, Zhifeng Chen, Yonghui Wu, and Macduff Hughes. Building machine translation systems for the next thousand languages, 2022.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Ka- mar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott M. Lundberg, Harsha Nori, Hamid Palangi, Marco Túlio Ribeiro, and Yi Zhang. Sparks of artificial general intelligence: Early experi- ments with GPT-4. CoRR, abs/2303.12711 2023. doi: 10448500arXi.22333.1777 URL attps:://oli wtt.:o.../.110011111... arxiv. 1..001111121
Isaac Caswell, Theresa Breiner, Daan van Esch, and Ankur Bapna. Language ID in the wild: Unex- pected challenges on the path to a thousand- language web text corpus. In Proceedings of the 28th International Conference on Computational Linguistics. pp. 6588 6608, Barcelona, Spain (On- line), December 2020. International Committee on Computational Linguistics. doi: 10.18653/ v1 wwt.00.coiin-mmin.555.. URL https::////... httl//acllaaaa wwwaaaaannoooooooooooooooooooooo ogy orgg2220. coling-maii 579.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2 03374, 2021.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. CoRR, abs/2110.14168. 2021. URL https:/ //trxaaaavvv///////////ee/aaaa org/ alos 22100111.... 14168.
Marta R. Costa-jusså, James Cross, Onur Çelebi, Maha Elbayad, Kenneth Heafill, Kevin Heffer- nan, Elahe Kalbassi, Janice Lam, Daniel Licht, Jean Maillard, Anna Sun, Skyler Wang, Guil- laume Wenzek, Al Youngblood, Bapi Akula, Loie Barrault, Gabriel Mejia Gonzalez, Prangthip Hansanti, John Hoffmman, Semarley Jarrett, Kaushik Ram Sadagopan, Dirk Rowe, Shannon Spruit, Chau Tran, Pierre Andrews, Necip Fazil Ayan, Shruti Bhosale, Sergey Edunov, Angela Fan, Cynthia Gao, Vedanuj Goswami, Francisco Guzmán, Philipp Kochn, Alexandre Mourachko, Christophe Ropers, Saflyyah Saleem, Holger Schwenk, and Jeff Wang. No language left
10

---

behind: Scaling human-centered machine translation. 10.488550/rx1v.2227..444 04672. URL https:: owt/10...1110..
CoRR, abs/2207.04672
2022. doi:
Siddharth Dalmia, Dmytro Okhonko, Mike Lewis, Sergey Edunov, Shinji Watanabe, Florian Metze, Luke Zeu,emoyer, and Abdelrahman Mohamed. LegoNN: Building Modular Deccoer Models, June 2022. URL ptctt:://axtie www//ogg/sss....oooo..... 2206. 03318. 8. arXiv:2206.033) 3818 [cs, eess].
Xavier Garcia, Aditya Siddhant, Orhan Firat, and Ankur P. Parikh. Harnessing multilinguality in unsupervised machine translation for rare languages. In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkami- Tür, Iz Beltagy. Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou (eds.), Proceedings of the 2021 Conference of the North Amer- ican Chapter of the Association for Computational Linguisties: Human Language Technolo- gies, NAACL-HLT 2021, Online, June 6-11, 2021, pp. 1126-1137. Association for Computa- tional Linguistics, 2021. doi: 10.18653/v120011 122655///2222ll.accmmmmaa mnaacl-main.8"" URL https:
10. 20.18653/v1/2000 www.nnaaaaaaaaaam..a 89.
Google, Rohan Anil, Andrew M. Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, Eric Chu, Jonathan H. Clark, Laurent EI Shafey, Yanping Huang, Katlorahy Meier-Hellsserm, Gaurav Mishra, Erica Moreira, Mark Omernick, Kevin Robinson, Sebastian Ruder, Yi Tay, Kefan Xiao, Yuanzhong Xu, Yujing Zhang. Gustavo Hernandez Abrego, Junwhan Ahn, Jacob Austin, Paul Barham, Jan Botha, James Bradbury, Siddhartha Brahma, Kevin Brooks, Michele Catasta, Yong Cheng, Colin Cherry, Christopher A. Choquette- Choo. Aakanksha Chowdhery, Clément Crepy, Shachi Dave, Mostafa Dehghani, Sunipa Dev, Jacob Devlin, Mark Díaz, Nan Du, Ethan Dyer, Vlad Feinberg, Fangxiaoyu Feng. Vlad Fienber, Markus Freitag. Xavier Garcia, Sebastian Gehrmann, Lucas Gonzalez, Guy Gur-Ari, Steven Hand, Hadi Hashemi, Le Hou, Joshua Howland, Andrea Hu, Jeffrey Hui, Jeremy Hurwitz, Michael Isard, Abe Ittycheriah, Matthew Jagielski, Wenhao Jia, Kathleen Kenealy, Maxim Krikun, Sneha Kudugunta, Chang Lan, Katherine Lee, Benjamin Lee, Eric Li, Music Li, Wei Li, YaGuang Li, Jian Li, Hyeontaek Lim, Hanzhao Lin, Zhongtao Liu, Frederick Liu, Marcello Maggioni, Aroma Mahendru, Joshua Maynez, Vedant Misra, Maysam Moussalem, Zachary Nado, John Nham, Eric Ni, Andrew Nystrom, Alicia Parrish, Marie Pel- lat, Martin Polacek, Alex Polozov, Reiner Pope, Siyuan Qiao, Emily Reif, Bryan Richter, Parker Riley, Alex Castro Ros, Aurko Roy, Brennan Sacta, Rajkumar Samuel, Renee Shelby, Ambrose Slone, Daniel Smilkov, David R. So, Daniel Sohn, Simon Tokumine, Dasha Valter, Vijay Vasude- van, Kiran Vodrahalli, Xuezhi Wang, Pidong Wang, Zirui Wang, Tao Wang, John Wieting, Yuhuai Wu, Kelvin Xu, Yunhan Xu, Linting Xue, Pengcheng Yin, Jiahui Yu, Qiao Zhang, Steven Zheng, Ce Zheng, Weikang Zhou, Denny Zhou, Slav Petrov, and Yonghui Wu. Palm 2 technical report, 2023.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin de Laroussilhe, An- drea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for NLP. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th Interna- tional Conference USA, volume 97 of Proceedings of Machine Learning Research, pp. 2790-2799. PMLR, 2019. URL hwt.:htt:/// ptwceedrooeeeeeeeoee ngs m1 www.pp/ssvvv.. hoouhhouleee houlsbyl 9a html.
Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California,
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Wang, and Weizhu Chen. LoRA: Low-Rank Adaptation of Large Language Models, October 2021. URL http:
http://ab///100 09685. arXiv:2106.099 00665 [cs].
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt,
Hannaneh Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. arXiv preprint arXiv:22/2.040089 2022.
Samuel Kessler, Bethan Thomas, and Salah Karout. An Adapter Based Pre-Training for Efficient and Scalable SuperSspe Speech Representation Learning, February 2022. URL http: 13530. arXiv:2107.13530 [cs, cess].
Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay V. Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-
11

---

Solo, Yuhuai Wu, Behnam Neyshabur, Guy Gur-Ari, and Vedant Misra. Solv- ing quantitative reasoning problems with language models. In NeurIPS, 2022. http://papers.ee nips wc/pape:pppcceeteeetc.cceee..eee www.:ffilsssppppeee wttappe///2//aaaaaaeeeaaa /aash/ .nw.errs.niiiss....
URL
html.
Shuai Lu, Daya Guo, Shuo Ren, Junjie Huang, Alexey Svyatkovskiy, Ambrosio Blanco, Colin B. Clement, Dawn Drain, Daxin Jiang, Duyu Tang, Ge Li, Lidong Zhou, Linjun Shou, Long Zhou, Michele Tufano, Ming Gong, Ming 7hou, Nan Duan, Neel Sundaresan, Shao Kun Deng, Shengyu Fu, and Shujie Liu. Codexglue: A machine learning benchmark dataset for code understanding and generation. CoRR, abs/2102.00664, 2021.
Jiaqi Ma, Zhe Zhao, Jilin Chen, Ang Li, Lichan Hong, and Ed H. Chi. SNR: sub-network routing for flexible parameter sharing in multi-task learning. In The Thirty-Third AAAI Con- ference on Artificial Intelligence, AAAI 2019, The Thirty-First Innovative Applications of Ar- tificial Intelligence Conference, IAAI 2019, The Ninth AAAI Symposium on Educational Ad- vances in Artificial Intelligence, EAAI 2019, Honolulu, Hawaii, USA, January 27 February 1, 2019, pp. 216 223. AAAI Press, 2019. doi: 10. /t000///aa..3311.3331.11.. //doi. org. 10. 1609/aaai.....
URL https:
v33101. 3301216.
Michael S Matena and Colin A Raffel. Merging models with fisher-weighted averaging. Advances in Neural Information Processing Systems, 11117703-1771,, 2022.
Mohammed Muqeeth, Haokun Liu, and Colin Raffel. Soft merging of experts with adaptive routing. arXiv preprint arXiv:2306.03745, 2023.
Jonas Pfeiffer, Aishwarya Kamath, Andreas Rücklé, Kyunghyun Cho, and Iryna Gurevych. Adapter- Fusion: Non-Destructive Task Composition for Transfer Learning, January 2021. URL http: / //arxiv.orgg org/ abs 2005 00247. arXiv:2005.00247 [cs].
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to- text transformer. J. Mach. Learn. Res., 21:140:1-140::: 2020. URL http:/ fjmrr. .errrorgg orgg
papers/v21/20-01 074 html.
Timo Schick, Jane Dwivedi-Yun Roberto Dessi, Roberta Raileanu, Maria Lomeli, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. Toolformer: Language models can teach themselves to use tools. CoRR, abs/2302.04761, 2023. doi: 1044550/ARXXV.2302.00777 //doi. org.1o........ 48550/ Aa00//rXXia aawii..iiiivv./iia 2302. 04761.
URL https:
Yongliang Shen, Kaitao Song, Xu Tan, Dongsheng Li, Weiming Lu, and Yueting Zhuang. Hug- gingGPT: Solving AI Tasks with ChatGPT and its Friends in HuggingFace, April 2023. URL http:://arxi org/ http/ab//23000 11.303111580. arXiv:2303. 17580 [es].
Freda Shi, Mirac Suzgun, Markus Freitag, Xuezhi Wang. Suraj Srivats, Soroush Vosoughi, Hyung Won Chung, Yi Tay. Sebastian Ruder, Denny Zhou, Dipanjan Das, and Jason Wei. Lan- guage models are multilingual chain-of-thought reasoners. In The Eleventh International Confer- ence on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https:
www.reveee.eeeeeeeeeepp.. httpss//
net http//dfficcc faafiid=fRwGGCk
Aditya Siddhant, Ankur Bapna, Orhan Firat, Yuan Cao, Mia Xu Chen, Isaac Caswell, and Xavier Garcia. Towards the next 1000 languages in multilingual machine translation: Exploring the synergy between supervised and self-supervised learning. CoRR, abs/2201.031 001100, 2022. URL https::/aarse
org/ abs 2201.22112222 03110.
Karan Singhal, Tao Tu, Juraj Gottweis, Rory Sayres, Ellery Wulezyn, Le Hou, Kevin Clark, Stephen Pfohl, Heather Cole-Lewis, Darlene Neal, Mike Schaekermann, Amy Wang, Mohamed Amin, Sami Lachgar, Philip Andrew Mansfield, Sushant Prakash, Bradley Green, Ewa Dominowska, Blaise Agüera y Arcas, Nenad Tomasev, Yun Liu, Renee Wong, Christopher Semturs, S. Sara Mahdavi, Joelle K. Barral, Dale R. Webster, Gregory S. Corrado, Yossi Matias, Shekoofeh Azizi, Alan Karthikesalingam, and Vivek Natarajan. Towards expert-level medical question answering with large language models. CoRR, abx/2305.096 2023, doi: 10.48550/arXiv.2300 096177. URL https:tt///// //tp:: wrooo1//000 48550/ axxi 200......0506666 09617.
12

---

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Anmual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp- 1998-6008, 2017. URL https: htt.tpps//ppccceciigg.soo new.eeiiiinggn..n.i....... http:/sweeriis....pparr/////0/ 3f5ee243547dee 100.777eee1110000111444448888 Abstract. html.
Ruize Wang, Duyu Tang. Nan Duan, Zhongyu Wei, Xuanjing Huang, Jianshu ji, Guihong Cao. Daxin Jiang, and Ming Zhou. K-Adapter: Infusing Knowledge into Pre- Pre-Trained Models with Adapters, December 2020. URL arXiv:2002.01800 [cs].
2b2/ 01808.
Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rehecca Roelofs, Raphael Gontijo Lopes, Ari S. Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. Model soups: averaging weights of multiple fine-tuned models im- proves accuracy without increasing inference time. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 23965 23998.. PMLR, 2022. URL https: html.
/ //proceeding nggs. ngs..... ml press / /vl 62
Prateek Yadav, Derek Tam, Leshem Choshen, Colin Raffel, and Mohit Bansal. Resolving interfer- ence when merging models. arXiv preprint arXiv:2306.0/708 2023.
Andy Zeng, Maria Attarian, Brian Ichter, Krzysztof Choromanski, Adrian Wong, Stefan Welker, Federico Tombari, Aveek Purohit, Michael Ryoo, Vikas Sindhwani, Johnny Lee, Vincent Van- houcke, and Pete Florence. Socratic Models: Composing Zero-Shot Multimodal Reasoning with Language, May 2022. URL [cs].
http://wwww....ra/////aa//////// 00598. arXiv:2204.00598
Zachary M. Ziegler, Luke Melas-Kyriazi, Sebastian Gehrmann, and Alexander M. Rush. Encoder- agnostic adaptation for conditional language generation. ArXiv, /et//908.069338 2019. URL https::/ api.. semanticschols ottg/Corruu
13

---

A SUPPLEMENTARY MATERIAL FOR NTL
A.1 FLORES-200
Figure depicts the gains the anchor PaLM2-S model when augmented with model that has been trained on DNTL. We see positive gain through CALM for 175 of 192 languages. The highest gains are seen for low-resource languages since they are the most underrepresented in the original model. Diminishing returns with higher resource languages is seen and this trend is similar to the trend seen for mg
15.0
Gain over Anohor Model
10.0
5.0
Low High Resource Languages (flanguages 192)
Figure 2: Gains seen by the composed model mл@B over the anchor model, mB, for the complete set of FLORES-200 languages. The languages are sorted from low to high-resource.
mA
mB
MAØB (CALM)
NTL B
MAØR (CALM)
B
NTL.
meo
7.6
28.8
34.0
33.2
bho
4.0
23.6
29.2
22.8
mfa
4.0
14.0
17.6
20.4
cv
6.0
17.6
16.4
20.4
pem
4.4
34.4
33.6
31.6
mni
3.6
2.8
4.4
6.0
efi
3.2
14.8
18.0
14.0
2.4
9.6
12.4
12.0
min
25.2
23.6
24.8
kri
5.6
12.4
18.8
20.0
ilo
4.8
14.8
16.8
14.0
tk
5.2
27.2
29.2
28.8
ady
6.4
30.0
36.4
29.2
gom
4.8
22.4
25.2
22.8
mai
3.2
22.8
24.8
21.2
ug
6.0
23.2
29.2
26.4
nso
6.0
8.4
8.4
9.6
ckb
3.2
25.6
28.0
27.2
4.8
31.6
36.4
27.6
1.2
5.2
9.2
4.0
bew
4.4
33.6
34.8
33.6
doi
3.6
17.2
22.4
21.6
4.8
7.2
10.0
11.6
dz
4.4
0.8
0.4
0.0
dv
2.8
11.2
14.8
13.2
avg.
4.5
18.6
21.4
19.8
Table 6: Performance evaluations on the complete set of low-resource languages GSM-8K. Augmenting ma with mg as m improves performance over mg across majority of languages. average, we see an improvement of 2.8%.
14

---

meo
mfa
pem
min
ilo
ady
Overlap
83.17
75.54
81.28
78.35
77.90
77.80
76.21
Delta
1.15
1.25
1.18
1.22
1.23
1.24
1.28
mai
nso
mzn
bew
ts
dv
bho
Overlap
76.63
69.58
71.32
71.37
61.62
55.18
73.67
Delta
1.26
1.40
1.38
1.37
1.55
1.70
1.30
cv
mni
or
kri
tk
gom
ug
Overlap
58.52
58.94
68.03
77.18
66.06
71.21
57.66
Delta
1.62
1.60
1.45
1.27
1.48
1.36
1.65
Table 7: Quality evaluation for the LRL GSM-8K dataset across languages. We created the dataset by translating the original English sentences of GSM-8K to the target language using the Google Translate API. We measure quality by back-translating the obtained examples back to English and measuring: (i) The overlap between the back-translated and the original English sentence, and (ii) The delta change in performance when PaI.M2-S evaluated on this sack-translated version of GSM-8K compared the original version.
A.2 GSM-8K
Quality evaluation for LRI. GSM-8K As described in Section 4.2, we created the GSM-8K dataset (Cobbe et al., 2021) for low-resource languages by using the Google Translate API to obtain silver translations in the target language from the source English sentence in the original dataset. We perform quality evaluation of these examples by back-translating them back to English using the same translation API defining two metrics over
(i) Overlap: The BLUE score measure between the actual example and the back trac-translated example, (ii) Delta: The change in performance of the PaLM2-S model when evaluated on the original GSM- 8K set compared to the back-translated version.
Table shows the values for these metrics across the various languages. We see that a decently high overlap value is seen across all languages. At the same time, the delta in performance is also minimal indicating that key attributes in the GSM-8K examples are not affected by translation.
Results on the complete language set Table 6 shows the comparative evaluations on the complete set 25 low-resource languages for which GSM evaluations are performed. We see an improvement over the anchor model mg for 20 of 25 languages. We also compare against the fully continued pre- trained version m
NTL and observe that MĄ@B outperform it for 18 of 25 languages.
B QUALITATIVE ANALYSIS
Table depicts a few qualitative examples for the code-to-text, or the code explanation task, for Python. These examples depict examples for the three broader bucket of examples that we observe in cases when CALM yields the correct responses:
When neither mA or ma generates the correct response but ma@r correctly attends over their latent representations to yield the correct output,
2. When either of mA or mB is seen to give the correct response while the other one is incor- rect and mA@B generates the correct response that matches the generation from the correct model of mA and my. and
3. When both mA and my generate the correct response and mAΦB reproduces those genera- tions.
We also observed similar qualitative patterns with other tasks for language inclusivity.
C OVERHEAD WITH CALM
In this section, we include a detailed computation of the expected parametric and training overhead while composing given models using our proposed CALM framework.
15

---

ConsumeBool (self)) try
value(self)
if aalf.has.value.
result ParseBool (self. token)
self. [OB]] get.val (K)
except ValueError as raise self. ParseError self. NextToken return result
else: raise ValueError found")
return
Consumes a boolean
Print error message and exit.
ma:
Consumes boolean
[a part of given model prefix]
The object is not a member
Exit with error message Print an error message and exit
CALM: Consumes boolcan
def get_positions (url): data -.et.resource (url)
def distance (x0, y0, x1, yl): return sqrt (pxx(xxxxxxx) pow pyw-y100,2)
positions [x for x data['p']]
return positions
Returns list of positions.
Returns the distance between two points
Positions of specified instruments.
Calculates the distance between two points
Get all positions.
Return the distance between points Calculates the distance between two points
Returns a list of positions.
Table Ccheryy-ickkre qualitative examples for the code-to-texs task on Python that depict examples that fall into a set of larger bucket of patterns that we observe across examples. CALM does well various settings: (i) when mxproduces the correct output but not mp, (ii) vice-versa- wwhen mp does well, and (iii) when neither of the two base models do well but combination of intermediate representations allow the composed model to give the correct output. This shows that composition implicitly learns do both: routing across models and combination, based given input.
16

---

C.I PARAMETRIC OVERHEAD
Building from the notations in $3.1, let's say the two models mA and ma have Na and NN number of standard transformer layers, respectively, with each layer of output dimensionality D and DB- As mentioned, we choose n = ng number of layers to perform the composition.
# params per fproj layer
(DA DB)
# params per Jenss layer
(3 D'B)
# total added params
Tx (DA DB +3 D)B)
# params in mg
NB (VB DB
+3 D3
+2* DB 3 * DB * KB)
where, Va and Kn depict the vocabulary size and hidden multiplication factor, respectively,
Let's consider some standard transformer configurations understand the parameter overhead. As an example, consider the layer configurations of standard BERT models: BERT-small (MA) and BERT-large (mg). That is: NA = 4, DA = 512. NB B = 24, Dg 1024, VB = 30K, KB = 4. Assuming that select all layers of ma, the value of 4. Then:
# params for CALM
= 4 (512 1024
+3 1024²)
23 1.5 x 10? & 15M
# params in mg
24 * (30K * 1024
+-3 10242
+2 10242 4) x 1B
%age added parameters
(15M/1B) * 100 1.5%
Hence, number of parameters added during composition is approximately 1.5% of those in ma.
C.2 TRAINING OVERHEAD
Although learning the new parameters during CALM training requires back propagating over mg. the total training costs are still significantly lesser than fine-tuning mg, owing the significantly lesser training examples.
Firstly, discussed in the previous section, the additional number of parameters introduced during composition is 11.5% of those in my hence negligible parametric addition. Further, since only ~5-7% of the total my fine-tuning data required train CALM, the training cost of CALM is minimal with respect to that of fine-tuning
Let's assume that (i) the number of parameters mB PB and the number of examples required to fine-tune mg is DB, (ii) the cost (in FLOPS), C(.), of training scales linearly with parameters and data, i.e., ((mg) = O(PB * DB) (iii) the number of parameters in mA is 10% of those in mg, (iv) the number of parameters added for CALM is 2% of those on mp, and (v) the amount of data required to train CALM is 5% my training. Then:
P PrRgE
PA PB Pec
= 0.10 PB + PB + 0.02 * Pg
1.12 Pp
DmHoE
= 0.05 DR
C(mA@B)
O(PmAce * Dmaaa)
O((1.12 PB) (0.05 Dn))
= O * PB * DB)
5.6% O(PB DB)
5.6% C(mB)
Hence, the cost of training CALM a significantly lesser (<10%) than that of fine-tuning mg.
17


Reviews:
[{"review_id": "zsCw4T0Ub9", "reviewer": "Reviewer_kQBq", "summary": "The paper proposes a new approach, CALM, to compose an anchor model and a domain-specific augmenting model by introducing a small number of trainable parameters over both models’ intermediate layer representations. This approach combines the capabilities of the anchor model and the augmenting model, and therefore, is able to address new challenging tasks that cannot be solved by either model alone. The authors use three tasks to demonstrate the advantages of CALM:\n1. solving arithmetic expressions containing keys: the anchor model is trained to do arithmetic over integers, and the augmenting model is trained to memorize string-to-integer key-value mappings.\n2. Translation and math-word problem-solving in low-resource languages: the anchor model is a pretrained PaLM-2 S model and the augmenting model is trained on low-resource languages.\n3. Code completion, text-to-code generation, and code-to-text generation: the anchor model is a pretrained PaLM-2 S model and the augmenting model is pretrained on codes.", "questions": "1. It is not clear to me how the projection function works. Does it project each selected layer of A to all selected layers of B? What is the definition of HBj? Does HA⊕Bj contain the information from all HA or only from a specific layer HAi?\n2. Why does CALM on the KV-Arithmetic task have a higher score than mB on the  Numeric-Arithmetic task? Should mB trained on the Numeric-Arithmetic task be the upper bound?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Strengths\n1. The new approach enables new capabilities in the composed version that the original models cannot achieve.\n2. CALM does not modify the parameters of the original model, avoiding the catastrophic forgetting that is prevalent in conventional approaches.\n3. CALM has a flexible architecture that makes it possible to compose more than one augmenting model with an anchor model.", "weaknesses": "1. The first task, key-value arithmetic, seems rather arbitrary. A simple encoder-decoder model is expected to solve this problem, leaving the readers the question of why they bother to use the composition of two models.\n2. For the second task, machine translation in low-resource languages, when the anchor model is trained on the low-resource languages, its performance is higher than the composed approach, which again raises the question of why you need to compose two models instead of fine-tuning the anchor model.\n3. The paper misses the details of parameterization/hyperparameter selection. For example, the authors did not write how the layers of the anchor and the augmenting models are selected.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nioxcZoSI5", "reviewer": "Reviewer_QMa2", "summary": "The paper introduces the CALM framework, which aims to enhance an existing LLM by incorporating a smaller, specialized one. This framework combines two pretrained LLMs while keeping their weights intact and utilizes learnable linear projections and self-attention to map intermediate representations from one model to the other. CALM demonstrates superior performance compared to both the anchor and augmenting models across three diverse tasks, including arithmetic reasoning, low-resource language translation, and code generation, which are not part of direct training but require a combination of the capabilities of both models.", "questions": "1.  Given that CALM requires the model weights access and running both forward and backward passes, keeping the weights frozen isn't necessarily required. It would be interesting to investigate whether updating an anchor, augmenting, or both models could enhance performance. \n2. It would be also interesting to explore the interactions between $m_B$ attending to $m_A$ in comparison to the current direction, to determine whether the capabilities of \"generic\" models can be effectively transferred to a \"specialized\" model, especially given that $m_A$ is more efficient to run.\n3. Table 3 shows more noticeable performance gains in high-resource languages compared to low-resource ones. I wonder if this might imply that the composition strengthens the skills of pretrained models thanks to extra parameters, rather than teaching them entirely new skills.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Comprehensive Experiments: The authors conduct extensive experiments that span three widely used application domains. The choices of training and evaluation tasks are reasonable and effectively show the composition of two LLMs.\n\n2. Performance Improvement: The experimental results consistently demonstrate performance improvements across all domains. Notably, the arithmetic reasoning experiment on synthetic datasets shows an impressive performance gain in comparison to both the anchor and augmenting models. The ability of CALM to tackle a novel task not previously feasible for either model is particularly intriguing.\n\n3. Usefulness: The paper presents a concept with potential real-world applications. It enhances model reusability and eliminates the need to scale up existing models for injecting new knowledge.\n\n4. Clarity: The paper is well-written and easy to follow.", "weaknesses": "1. Lack of Comparison to Existing Methods: While the idea of combining models with different skill sets has been explored in related work, such as ensembling hidden representations of two models, the paper lacks empirical results for such comparisons. Authors are recommended to consider providing a comparison of CALM with existing methods for ensembling model representations. Additionally, could the authors clarify \"LoRA assumes access to the full underlying domain data during pretraining\"? It is assumed that the same training data are used for both LoRA and CALM in Table 5.\n\n2. Missing Information on Overhead: The paper does not provide information on the additional parameters introduced and their ratio to the anchor and augmenting models, which is important to understand the system's overhead. It would also be helpful to support the claim that \"The composed model achieves similar performance for a tiny fraction of the training cost\" with evidence, such as the number of FLOPs and memory requirements. Since the framework requires backpropagation through the initial layers of both models, I assume that the training cost should be comparable to directly training $m_B$.\n\n3. The paper should include model configurations and training details to enhance reproducibility.\n\n4. Some notations are used before being formally introduced, leading to potential confusion. For instance, notation $\\mathbf{C}$ is introduced in Section 3.2 after being used from the beginning of Section 3. Additionally, both $t_{\\{A, B\\}}$ and $C$ denote task sets, which can be confusing and should be clarified or distinguished.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cZPYhuBzXO", "reviewer": "Reviewer_6aS2", "summary": "The paper proposes CALM, composition to augment language models. CALM is designed to compose LLMs with different capabilities. In particular, the paper focuses on composing an anchor model with a domain-specific augmenting model to enable new capabilities. CALM introduces simple linear transformations to extract features from the augmenting model and insert them into the anchor model with cross-attention modules. In the experiments, the paper introduces three different applications of CALM, in the domains of key-value arithmetic, low-resource language inclusivity, and code understanding and generation. The experimental results demonstrate the composed LLMs can perform new challenging tasks and obtain better results than the models without composition.", "questions": "Many studies have explored how to efficiently connect image encoders to LLMs, which finally achieves a \"combined skill\" like image-to-text generation. What is the difference between CALM and the efficient cross-modal LLM methods?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The idea of LLM composition is interesting. A great number of LLMs are designed with diverse capabilities. Combining LLMs with different capabilities and deriving new capabilities are valuable research questions. In section 4.2, after combining the model with KV-substitution skill and the model with numeric-arithmetic skill, the combined model achieves zero-shot KV-arithmetic inference. In comparison, the underlying models fail to do KV-Arithmetic inference, demonstrating the emergence of \"new skills\".\n- The authors conduct extensive experiments on diverse domains ranging from language inclusivity to code generation. The results show that CALM successfully combines the capabilities of the anchor and augmenting models.", "weaknesses": "- Except for the key-value arithmetic case, the composition between low-resource languages and English, and the composition of coding and language capabilities are similar to the works of efficient cross-modal LLMs. For example, many studies have explored how to efficiently connect image encoders to LLMs, which finally achieves a \"combined skill\" like image-grounded language generation. I believe the relation between CALM and these related cross-modal models should be discussed. Besides, I would expect the combined models to have more non-trivial combined skills ( key-value arithmetic) beyond conditional generation (machine translation or code-to-text generation).\n- Introduction mentions that directly training LLMs is computationally expensive but it seems that the experiments do not involve whether CALM is more computationally efficient than directly training the anchor model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xMdkPuM7bS", "reviewer": "Reviewer_AR8D", "summary": "This paper proposes CALM to compose a general LLM with specialized models with a cross-attention module. By only fine-tuning the cross-attention on a small amount of data, the composed system is able to enable new capabilities for compositional tasks that neither of the two models can handle independently. The authors conduct experiments on tasks like math reasoning on low-resource language and code understanding and generation to verify the effectiveness of the proposed method.", "questions": "Please see above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Though similar ideas are tested to be effective for building multi-modal models (e..g. reuse image encoder outputs for an LLM to enable multi-modal capabilities), this paper extends the application to new domains and scenarios.\n- The method is simple and effective.", "weaknesses": "- Lack of important details, e.g. it is unclear how many the size of these models, the training steps, and hyperparameters. And how large is the training data $D_c$?\n- It would be helpful to add a discussion on the composition with existing methods, such as routing and tool-using. If we treat the specialized models as external tools and prompt the general LLM to first call these models and then process the returned outputs, is it a kind of composition in the text space rather than representations?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)