Communicating Activations Between Language Model Agents — part 2/3

Cleaned-text code points [14463, 29698). Bibliography omitted here; source.md includes it.

Let 𝒂 = ( 𝒉 A , k ) t A \bm{a}=(\bm{h}_{A,k})_{t_{A}} , 𝒃 = ( 𝒉 B , j ) t B \bm{b}=(\bm{h}_{B,j})_{t_{B}} . For sake of simplicity assume d A = d B d_{A}=d_{B} . 1 1 1 When d A ≠ d B d_{A}\neq d_{B} , the 𝚜𝚞𝚖 \mathtt{sum} , 𝚖𝚎𝚊𝚗 \mathtt{mean} , and 𝚛𝚎𝚙𝚕𝚊𝚌𝚎 \mathtt{replace} functions are defined as follows. Let d = min ⁡ ( d A , d B ) d=\min(d_{A},d_{B}) and ∘ \circ the concatenation operator. Then f ( 𝒂 , 𝒃 ) = 𝒃 1 : max ⁡ ( d B − d , 0 ) ∘ ( 𝒃 max ⁡ ( d B − d , 0 ) + 1 : d B + 𝒂 max ⁡ ( d A − d , 0 ) + 1 : d A ) f(\bm{a},\bm{b})=\bm{b}_{1:\max(d_{B}-d,0)}\circ\left(\bm{b}_{\max(d_{B}-d,0)+1:d_{B}}+\bm{a}_{\max(d_{A}-d,0)+1:d_{A}}\right) ( 𝚜𝚞𝚖 ) \mathtt{(sum)} , f ( 𝒂 , 𝒃 ) = 𝒃 1 : max ⁡ ( d B − d , 0 ) ∘ 1 2 ( 𝒃 max ⁡ ( d B − d , 0 ) + 1 : d B + 𝒂 max ⁡ ( d A − d , 0 ) + 1 : d A ) f(\bm{a},\bm{b})=\bm{b}_{1:\max(d_{B}-d,0)}\circ\frac{1}{2}\left(\bm{b}_{\max(d_{B}-d,0)+1:d_{B}}+\bm{a}_{\max(d_{A}-d,0)+1:d_{A}}\right) ( 𝚖𝚎𝚊𝚗 ) \mathtt{(mean)} , and f ( 𝒂 , 𝒃 ) = 𝒃 1 : max ⁡ ( d B − d , 0 ) ∘ 𝒂 max ⁡ ( d A − d , 0 ) + 1 : d A f(\bm{a},\bm{b})=\bm{b}_{1:\max(d_{B}-d,0)}\circ\bm{a}_{\max(d_{A}-d,0)+1:d_{A}} ( 𝚛𝚎𝚙𝚕𝚊𝚌𝚎 ) \mathtt{(replace)} . We consider three non-learned functions f f : f ⁡ ( 𝒂 , 𝒃 ) \displaystyle f(\bm{a},\bm{b}) = 𝒂 + 𝒃 ( 𝚜𝚞𝚖 ) \displaystyle=\bm{a}+\bm{b}\>\>\qquad\qquad\mathtt{(sum)} f ⁡ ( 𝒂 , 𝒃 ) \displaystyle f(\bm{a},\bm{b}) = 1 2 ( 𝒂 + 𝒃 ) ( 𝚖𝚎𝚊𝚗 ) \displaystyle=\frac{1}{2}(\bm{a}+\bm{b})\>\>\qquad\mathtt{(mean)} f ⁡ ( 𝒂 , 𝒃 ) \displaystyle\qquad f(\bm{a},\bm{b}) = 𝒂 ( 𝚛𝚎𝚙𝚕𝚊𝚌𝚎 ) \displaystyle=\bm{a}\>\>\qquad\qquad\mathtt{(replace)}

For cases where, due to differences in A A and B B ’s training, A A and B B ’s activation spaces are quite different, we propose learning a task-agnostic (depends only on the models A A and B B ) linear layer 𝑾 ∈ ℝ d B × ℝ d A \bm{W}\in\mathbb{R}^{d_{B}}\times\mathbb{R}^{d_{A}} that projects 𝒂 \bm{a} onto B B ’s activation space. Note that this introduces zero additional task-specific parameters and data, as we propose learning this “mapping matrix” 𝑾 \bm{W} only once for each model pair ( A , B ) (A,B) using general text, e.g. sequences from A A and/or B B ’s pretraining data mixes. We can then perform 𝚜𝚞𝚖 \mathtt{sum} , 𝚖𝚎𝚊𝚗 \mathtt{mean} , or 𝚛𝚎𝚙𝚕𝚊𝚌𝚎 \mathtt{replace} with 𝑾 ​ 𝒂 , 𝒃 \bm{W}\bm{a},\bm{b} instead of 𝒂 , 𝒃 \bm{a},\bm{b} . We propose training 𝑾 \bm{W} to minimize MSE loss over a dataset of N N sentences ℒ MSE ​ ( { 𝒚 ( i ) } i = 1 N , { 𝒛 ( i ) } i = 1 N ) = 1 N ​ ∑ i = 1 N ‖ 𝒛 ( i ) − 𝑾 ​ 𝒚 ( i ) ‖ 2 2 \mathcal{L}_{\rm MSE}\left(\{\bm{y}^{(i)}\}_{i=1}^{N},\{\bm{z}^{(i)}\}_{i=1}^{N}\right)=\frac{1}{N}\sum_{i=1}^{N}\left\|\bm{z}^{(i)}-\bm{W}\bm{y}^{(i)}\right\|_{2}^{2} where each ( 𝒚 ( i ) , 𝒛 ( i ) ) (\bm{y}^{(i)},\bm{z}^{(i)}) pair denotes the final-token layer- 26 26 activations of A A and B B at layers k k and j j respectively given the same sentence as input.

### 3.2 Compute Analysis

To understand the significance of activation communication, we must formally quantify the compute this procedure saves over natural language communication. For simplicity suppose the following (similar calculations can be made for the cases where A A and B B have differing model architectures and/or are given different prompts): • A A and B B both have L L layers (each with H H attention heads, key size K K , and feedforward size F F ), dimension D D , and vocab size V V

• A A and B B are both given a prompt of P P tokens

• A A can send B B a single M M -token message

• B B must produce an output of T T tokens, given its prompt and A A ’s message

Traditional methods require M M forward passes of A A given a P P -length input, plus T T forward passes of B B given a ( P + M ) (P+M) -length input. Following Hoffmann et al. (2022) , this requires M ⁡ ( 4 ​ P ​ V ​ D + L ⁡ ( 8 ​ P ​ D ​ K ​ H + 4 ​ P 2 ​ K ​ H + 3 ​ H ​ P 2 CLOSE CLOSE \displaystyle M\big(4PVD+L(8PDKH+4P^{2}KH+3HP^{2} (1) OPEN OPEN + 4 ​ P ​ D ​ F ) ) + T ⁡ ( 4 ​ ( P + M ) ​ V ​ D + L ⁡ ( 8 ​ ( P + M ) ​ D ​ K ​ H CLOSE CLOSE \displaystyle+4PDF)\big)+T\big(4(P+M)VD+L(8(P+M)DKH OPEN OPEN + 4 ​ ( P + M ) 2 ​ K ​ H + 3 ​ H ​ ( P + M ) 2 + 4 ​ ( P + M ) ​ D ​ F ) ) \displaystyle+4(P+M)^{2}KH+3H(P+M)^{2}+4(P+M)DF)\big) FLOPs. In contrast, at inference time, our method requires only 1 partial (up till the k k th layer) forward pass of A A given a P P -length input, T T forward passes of B B given a P P -length input, and the activation replacement procedure. This requires 2 ​ P ​ V ​ D + k ⁡ ( 8 ​ P ​ D ​ K ​ H + 4 ​ P 2 ​ K ​ H + 3 ​ H ​ P 2 CLOSE \displaystyle 2PVD+k(8PDKH+4P^{2}KH+3HP^{2} (2) OPEN + 4 ​ P ​ D ​ F ) + T ⁡ ( 4 ​ P ​ V ​ D + L ⁡ ( 8 ​ P ​ D ​ K ​ H + 4 ​ P 2 ​ K ​ H CLOSE CLOSE \displaystyle+4PDF)+T\big(4PVD+L(8PDKH+4P^{2}KH OPEN OPEN + 3 ​ H ​ P 2 + 4 ​ P ​ D ​ F ) ) + ℱ ⁡ ( D ) \displaystyle+3HP^{2}+4PDF)\big)+\mathcal{F}(D) FLOPs, where ℱ ⁡ ( D ) = O ⁡ ( D ) \mathcal{F}(D)=O(D) for non-learned f f and O ⁡ ( D 2 ) O(D^{2}) when f f is the mapping matrix.

In all practical cases, (2) is substantially lower than (1).

### 3.3 Why should this work?

Recall that Pham et al. (2024) propose CIPHER—communicating the average tokenizer embedding (weighted by the LLM’s next-token probabilities) between models. We build upon the intuition behind CIPHER, which goes as follows: the token sampling process during decoding risks substantial information loss from the model’s output logits, and communicating a model’s weighted-average tokenizer embedding essentially entails communicating both that model’s final answer and its belief in that answer (over the entire vocabulary).

Communicating activations, then, can be thought of as communicating a strict superset of {next-token prediction, belief over entire vocabulary}, as activations of late-enough layers essentially encode the model’s entire knowledge about the provided context as well as its predicted completion and confidence in that completion (see Figures 1 and 7 in Hewitt & Manning (2019) and Hernandez et al. (2024) , respectively, which show that linear probes tasked with predicting certain output characteristics from a Transformer’s intermediate layer embeddings of its input work poorly for early layers, extremely well after around the halfway point of computation, but then probe accuracy drops closer to the final layers). 2 2 2 Note one important critique of multiagent debate: that in cases where multiple agents are uncertain about the answer, there is no reason why referencing other agents’ answers would generate more factual reasoning. Both CIPHER and activation communication solve this problem, as some notion of model confidence is being communicated along with its next-token prediction. Indeed, these curves of probe accuracy by layer indicate that the final layers and LM head “ throw away ” information not useful for next-token prediction that very well could be useful for communicative purposes; this is precisely why our proposed activation communication technique is not an iterative approach (there is no notion of “rounds” like in debate and CIPHER, which require an additional token budget to extract more and more information out of the LM), as one activation grafting step from A A to B B inherently communicates to B B all of A A ’s knowledge/beliefs about the prompt it was given. Moreover, the extra information over the model’s next-token prediction and confidence that is encoded in its activations is what makes activation communication more performant than its natural language counterpart, as we will see in Section 4 .

## 4 Experiments

We test our method on two distinct experimental setups: multi-player coordination games ( Section 4.1 ) and reasoning benchmarks ( Section 4.2 ). Qualitative results are available in Appendix A .

### 4.1 Multi-player coordination games

Drawing from existing literature on multi-agent communication, we design two Lewis signaling games ( Lewis, 2008 ; Lazaridou et al., 2016 ) to test the efficacy of activation communication (example prompts and answers in Table 1 ): 1. Countries , where A A is given as input a string of the format “ [ 𝙿𝙴𝚁𝚂𝙾𝙽 ] 𝚒𝚜 𝚊𝚝 𝚝𝚑𝚎 [ 𝙻𝙰𝙽𝙳𝙼𝙰𝚁𝙺 ] \mathtt{[PERSON]\text{ }is\text{ }at\text{ }the\text{ }[LANDMARK]} ” and B B is asked “ 𝚆𝚑𝚒𝚌𝚑 𝚌𝚘𝚞𝚗𝚝𝚛𝚢 𝚒𝚜 [ 𝙿𝙴𝚁𝚂𝙾𝙽 ] 𝚕𝚘𝚌𝚊𝚝𝚎𝚍 𝚒𝚗 ? \mathtt{Which\text{ }country\text{ }is\text{ }[PERSON]\text{ }located\text{ }in?} ”

2. Tip Sheets (inspired by Lewis et al. (2017) ), where A A is given a simulated “tip sheet” and B B is asked to make an informed investment decision in accordance with the information in the tip sheet.

We synthetically generate 100 100 ( Countries ) and 70 70 ( Tip Sheets ) different prompts and answers of the same format as the samples in Table 1 , and report the proportion out of those samples that B B responds with an exact string match to the ground truth answer. As baselines, we consider a “silent” (✗) setup, where the agents are not allowed to communicate; a “single-agent skyline,” where a single LLM is given the concatenation of A A and B B ’s prompts; and traditional natural language communication, where A A is asked to output a message that is then given to B B along with x B x_{B} . All decoding is done greedily.

Table 2 presents the results for both coordination games using 2 different instances of the same model as the agents ( A = B A=B ). Across the 3B and 8B model sizes, activation communication (AC) with f = 𝚛𝚎𝚙𝚕𝚊𝚌𝚎 f=\mathtt{replace} almost completely recovers the gap between the zero-communication (✗) and the single-agent skyline ( Skyline ), outperforming natural language communication (NL) using far less compute. We hypothesize that 𝚛𝚎𝚙𝚕𝚊𝚌𝚎 \mathtt{replace} is more effective than 𝚖𝚎𝚊𝚗 \mathtt{mean} and 𝚜𝚞𝚖 \mathtt{sum} as the former is guaranteed to output a vector within B B ’s activation space, while the latter two likely do not (e.g., the norm of the vector outputted by 𝚜𝚞𝚖 \mathtt{sum} will be around double that of a typical activation). Furthermore, most of the information B B needs is likely contained in its representations of previous tokens in the sequence, hence losing its final-token representation does not hurt.

### 4.2 Reasoning Benchmarks

Next, we test our methods on a variety of reasoning benchmarks, spanning several real-world tasks and domains.

#### Baselines

We benchmark activation communication against the following two baselines:

• Single Model : A single LLM responds to the prompt in natural language.

• Natural Language Debate (NLD) ( Du et al., 2023 ) : Each LLM provides an initial response to the given prompt. Then, for each of r − 1 r-1 subsequent rounds, each LLM is prompted to refine its previous response given the other agents’ responses as input. Note that NLD is the most direct baseline for our approach, as it is a state-of-the-art natural language communication protocol. We fix r = 2 r=2 in our experiments.

Note that we do not compare to Pham et al. (2024) , as they communicate the input (tokenizer) embeddings rather than activations/output embeddings between models, and hence require a shared tokenizer and embedding table between agents which is extremely restrictive and prevents applicability to our experimental setup.

To determine the values of k k and j j for activation communication (AC), we compute the accuracy on Countries and Tip Sheets for every pair ( k , j ) ∈ { 1 , … , 30 } 2 (k,j)\in\{1,\dots,30\}^{2} . Based on these results (shown in Figure 2 ) as well as Table 2 , we fix k = j = 26 k=j=26 and f f = = 𝚛𝚎𝚙𝚕𝚊𝚌𝚎 \mathtt{replace} for the following experiments.

Across all experiment configurations, we fix the decoding strategy to nucleus sampling with p = 0.9 p=0.9 .

#### Models

We conduct most of our experiments using 𝙻𝙻𝚊𝙼𝙰 \mathtt{LLaMA} - 3.2 \mathtt{3.2} - 𝟹 ​ 𝙱 \mathtt{3B} and 𝙻𝙻𝚊𝙼𝙰 \mathtt{LLaMA} - 3.1 \mathtt{3.1} - 𝟾 ​ 𝙱 \mathtt{8B} as the two agents. Additionally, to test our approach’s robustness and generalizability, we conduct experiments with models belonging to various other suites within the 𝙻𝙻𝚊𝙼𝙰 \mathtt{LLaMA} family and of several different sizes.

Note that for these experiments, we restrict the setting to communication between different models (rather than multiple instances of the same model in Section 4.1 ), since the same model would have identical activations for the same prompts, meaning no information would be communicated in the grafting process. We argue that the multiple-model setting is realistic (perhaps more so than the setting of multiple instances of the same model), as recent advances in LLM development have led to the release of models with specialized abilities ( Singhal et al., 2023 ) and of different sizes ( Dubey et al., 2024 ) that merit complementary usage. Our work thus answers the question: How can we get the best performance by leveraging multiple models of distinct capabilities and sizes, relative to the added inference-time compute over a single forward pass through any single model?

#### Datasets

We evaluate our technique on seven reasoning datasets that span various real-world tasks and domains: (i) Biographies ( Du et al., 2023 ) , which asks the LLM to generate a factual biography of a famous computer scientist; (ii) GSM8k ( Cobbe et al., 2021 ) , a variety of grade school math problems created by human problem writers; and (iii) 5 datasets randomly drawn from MMLU ( Hendrycks et al., 2021 ) : High School Psychology (from the Social Sciences category), Formal Logic (from the Humanities category), College Biology (from the STEM category), Professional Law (from the Humanities Category), and Public Relations (from the Social Sciences category). We evaluate on a randomly-sampled size- 100 100 subset of each dataset.

In experiments involving the mapping matrix 𝑾 \bm{W} , we instantiate 𝑾 ∈ ℝ 4096 × 3072 \bm{W}\in\mathbb{R}^{4096\times 3072} using Xavier initialization and train for 10 10 epochs on a dataset of 3072 3072 sentences 3 3 3 We use 3072 3072 sentences as linear regression with d d -dimensional input has a sample complexity of O ⁡ ( d ) O(d) ( Vapnik, 1999 ) . randomly drawn from the Colossal Clean Crawled Corpus (C4) ( Dodge et al., 2021 ) . We use batch size 32 32 and the Adam optimizer with learning rate 0.001 0.001 .

#### Metrics

We measure the accuracy of the final response for the single models and AC. For NLD, we measure the accuracy of the majority-held final-round answer across agents when the answer is automatically verifiable (numeric in GSM8k, multiple choice for the MMLU datasets) or the average final-round answer across agents otherwise (Biographies).

For GSM8k and the MMLU datasets, we report the proportion of samples in the dataset for which the generated answer exactly matches the ground-truth answer. For Biographies, following Du et al. (2023) , we prompt an LLM judge ( 𝙻𝙻𝚊𝙼𝙰 \mathtt{LLaMA} - 3.1 \mathtt{3.1} - 𝟾 ​ 𝙱 \mathtt{8B} ) to check whether each manually-decomposed fact in a ground-truth biography is supported ( 1 1 ), partially supported ( 0.5 0.5 ), or unsupported ( 0 0 ) in the generated biography, taking the mean of these scores over all facts as the per-biography accuracy and the mean over all dataset samples as the total accuracy.

#### Comprehensive evaluation with the 𝙻𝙻𝚊𝙼𝙰 \mathtt{LLaMA} family