Communicating Activations Between Language Model Agents

We thank the reviewer for their feedback, and address their comments below:

1. The theoretical depth of this paper can be further enhanced. For example, the result shows that the direct replacement strategy is the most efficient way. Basically, it means that we should discard all previous representations and directly use the projected ones. Do we have any assumptions behind this phenomenon? (e.g., what is the relationship between the two models such that this conclusion holds?)

First, note that the direct replacement strategy only replaces $B$'s layer-$j$ activation of the final token with $A$'s last-token layer-$k$ activation. The embeddings at all other token positions remain the same. Hence, we are not actually discarding all previous representations: after applying masked attention in each of the previous Transformer layers, the last token activation of $A$ attends to all tokens before it, hence incorporating information from the entire sequence; and the previous token activations of $B$ are retained and this incorporate all of $B$'s "thoughts" regarding the full sequence.

Please refer to Appendix B.1 for additional discussion and comparison between the replacement, sum, and mean strategies for activation grafting.

2. What is the connection and difference of your model against the encoder-decoder model?

This is an interesting point; the direct replacement activation grafting approach could be seen as running an encoder-decoder Transformer, where the first $k$ layers of $A$ are the encoder, $B$ is the decoder, and we have a "special cross-attention" where----from layer $j$ onwards----all tokens in $B$ attend to only the last-token embedding outputted by the first $k$ layers of $A$. However $A$ and $B$ are language models used out-of-the-box, with no joint training, no requirement of shared tokenizer/aligned embedding space/similar training distributions/etc.

3. How to choose the hyperparameters of j,k

Please see point #3 in the reply to Reviewer VKxi above.

4. What is the effect of the scale and distribution of the aligning data for training the transformation matrix?

We propose training $\mathbf{W}$ to minimize MSE loss over a dataset of $N$ sentences $\frac{1}{N}\sum_{i=1}^N \left\|\mathbf{z}^{(i)} - \mathbf{W}\mathbf{y}^{(i)}\right\|_2^2$ where each $(\mathbf{y}^{(i)},\mathbf{z}^{(i)})$ pair denotes the final-token layer-$26$ activations of $A$ and $B$ at layers $k$ and $j$ respectively given the same sentence as input.

Hence, the "aligning data distributions" are quite literally $A$ and $B$'s activation distributions; we do not scale or otherwise modify the activation vectors in any way before training.

5. Are you trying to train a unique matrix for each j,k pair? Since you are using the MSE loss, I wonder if such a training objective might work well given that the capacity of a matrix is not very big.

Yes we are, see point #4 above. However, note that for each model pair, we only need a single $(j,k)$ pair----indeed, we attain SOTA results with $k=j=26$ fixed.

Also, this training objective is quite standard in related literature. For instance, the "model stitching" paper [1] learns a 1x1 convolution between activations of two models at a specific layer, which is exactly equivalent to our method of training $\mathbf{W}$; [2] also learns a linear projection between Transformer layers which is shown to be quite expressive.

Indeed, we find our approach yields quite strong results; please see Appendix B.3 for additional discussion.

6. Have you tried directly aligning the representations from two models and adding one more layer on top of that?

We kindly ask for clarification on this question. If by "directly aligning the representations from two models and adding one more layer on top of that" you mean learning a linear layer that projects activations from $A$'s activation space to $B$'s, that is exactly what we do in the AC ($\mathbf{W}$) method.

We hope that with the provided responses & additional strong results, you consider raising your score to support clear acceptance. Please let us know if there are any additional questions or concerns; we'd be happy to address them.

[1] Yamini Bansal, Preetum Nakkiran, Boaz Barak. Revisiting Model Stitching to Compare Neural Representations, 2021.

[2] Alexander Yom Din et al. Jump to Conclusions: Short-Cutting Transformers With Linear Transformations, 2024.

We thank the reviewer for their feedback, and address their comments below:

1. All tested models come from the LLaMA family...

Please see point #1 in the reply to Reviewer VKxi above. In summary, we test AC using models across the LLaMA, Qwen, and Gemma families, and find that AC beats NLD across the board, and beats both individual models for 4/5 of the 6 model pairs on Biographies/GSM8k respectively----demonstrating the efficacy of AC irrespective of model architecture, size, tokenizer, and training data.

2. The paper lacks an analysis of why the method work, e.g. inspecting the geometric properties...

First, please see Section 3.3 for a theoretical grounding of this work.

Second, we conduct the following experiment: for each of the 6 pairs of models $A,B$ in the above experiment (see the table from response to Reviewer VKxi), we compute the increase in Biographies performance with AC relative to the average individual performance of $A$ and $B$. We also compute the matrix analog of squared cosine similarity between the models' activation spaces, $||Y^T X||_F^2/(||X||_F^2||Y||_F^2)$ where $X$ is the matrix of $A$'s activations on 3072 sentences from C4 (same dataset used to train $\mathbf{W}$), $Y$ is the same for $B$, and $||\cdot||_F$ is the Frobenius norm. This gives us the following plot (please click link if image not displayed):

https://i.ibb.co/3y6XF7Z7/model-comparison-plot.png

There is a clear positive correlation between similarity of the activation distributions and AC performance gain, as expected; the more aligned $A/B$'s activation spaces, the more semantically meaningful/useful the embedding we graft from $A$ to $B$.

3. Though not fundamental to position the paper, broader comparisons...

Thanks for sharing this. We've included a discussion of this and similar papers (e.g., Kornblith19, Similarity of Neural Network Representations Revisited) in the paper.

4. A link that I didn't see in the paper...

We thank the reviewer for bringing this to our attention. The "tuned lens" and "logit lens" [1] are related to the theoretical intuition behind this approach that we share in Section 3.3; we have added a discussion of these papers to both Section 3.3 and related works.

5. In the appendix I found...

Thanks for pointing this out, we have ensured all experiments in the appendix are either moved to or discussed in the main body of the paper.

6. NLD debate is found to improve with increased number of agents...

First, we want to distinguish between "agents" (distinct model instances) and "models" (distinct LMs).

With debate, more agents can help because we are sampling different outputs from each agent (model instance), and this can yield diverse reasoning paths that are recombined to produce stronger final outputs.

This isn't true with AC, as one activation grafting step from $A$ to $B$ inherently communicates to $B$ all of $A$'s knowledge/beliefs about the prompt it was given. We argue this is actually a benefit of AC over NLD, as we don't require increasing token budgets to extract more and more information out of the LMs.

A similar argument can be made for the number of rounds in NLD. Indeed, as shown in Appendix B.5, for 5 of the 7 reasoning benchmarks, AC beats NLD even with 3-4 rounds while using substantially less compute, highlighting the superiority and robustness of activations as an alternative “language” for inter-LM communication.

AC could theoretically scale to more than 2 models. While letting $f$ be the average function would work out of the box, we saw that direct replacement is much more effective. So, for instance, consider a setup (using the notation from Section 3.1 of our paper) where for any $i, j, k, m$ with $j < k$, we:

1. run a partial forward pass $B_{\leq j}(x_B)$ to get last-token activation $\mathbf{b}_j$;
2. run a partial forward pass $A_{\leq i}(x_A)$ to get $\mathbf{a}_i$;
3. replace $\mathbf{b}_j \leftarrow f(\mathbf{a}_i, \mathbf{b}_j)$;
4. continue $B$'s forward pass till layer $k$ to get last-token activation $\mathbf{b}_k$;
5. run a partial forward pass $C_{\leq m}(x_C)$ to get $\mathbf{c}_m$;
6. replace $\mathbf{c}_m \leftarrow f(\mathbf{b}_k, \mathbf{c}_m)$;
7. then continue $C$'s forward pass till decoding is complete.

This 3-model setup can extend to an arbitrary number of models. We leave this extension of our approach to future work.

7. The paper would be a clear accept with tests on a broader range of model families to confirm that activation communication is not specific to Llama's architecture or training paradigms, and additional insights on the effectiveness of the proposed approach.

We hope that with the provided responses & additional strong results, you consider raising your score to support clear acceptance. Please let us know if there are any additional questions or concerns; we'd be happy to address them.

[1] nostalgebraist. Interpreting gpt: the logit lens, 2020.

We thank the reviewer for their feedback, and address their comments below:

1. Lack of comparison with other similar approaches such as model merging

To adequately scope our paper, we chose to limit our focus to task-agnostic methods. Model composition/merging methods are extensively discussed in Sections 2 and 3.2, but they require a substantial amount of additional task-specific parameters and data, hence we do not compare to these. Furthermore these methods require much more compute in the form of LM finetuning, layer or router training, etc.; our approach is far more compute efficient.

2. The approach needs access to the models’ parameters which is not feasible for state-of-the-art close-source LMs.

Exploring API-only approaches is highly limiting (see AC's meta review of [1] at ICLR last year). Furthermore, recent releases of powerful open-source models [2] merit the development of embedding-based techniques.

3. Need for model-model pair training of W to map representations from one LM to another LM if architecture is different.

Note that learning $\mathbf{W}$...

• needs to happen exactly once for each model pair
• introduces zero additional task-specific parameters and data by virtue of requiring only general text, e.g. sequences from $A$ and/or $B$’s pretraining data mixes
• is quite sample-efficient: as mentioned in the paper, we use $3072$ sentences to train $W\in \mathbb{R}^{4096\times 3072}$, since linear regression with $d$-dimensional input has a sample complexity of $O(d)$ [3]

Furthermore, we find empirically that even when models have different architectures, tokenization schemes, or training distributions, we do not need to train a mapping matrix $\mathbf{W}$ to attain SOTA results. Please see point #1 in the reply to Reviewer VKxi above.

4. The work by Phametal.,2024, though different, should also be a method to compare as a baseline to better understand the effect of using intermediate activations versus embeddings in multi-agent frameworks.

Pham24 propose communicating the input (tokenizer) embeddings between models, meaning the two models must have the same tokenizer and embedding table to even run their approach. This severely limits the applicability of their method; in particular, given that all our experiments use model pairs (e.g., LLaMA-3.2-3B and LLaMA-3.1-8B) with distinct tokenizers and/or embedding layers, we unfortunately cannot compare against Pham24. Also, note Pham24's approach still faces substantial information loss relative to the model activations and, more importantly, does not save compute, as the number of embeddings passed between models is the same as the number of tokens passed between models in natural language communication.

5. There are also highly similar approaches that also rely on using intermediate activations from a language model to inject into intermediate layers of the same language model to improve reasoning (e.g. https://arxiv.org/abs/2412.06769 and https://arxiv.org/abs/2502.05171). These approaches do not need another training stage to train a separate W to map activations from one LM to another. More discussion about these similar works would help the reader better understand the difference.

Thank you for raising this point; we have added discussion of these approaches to our paper. In summary, such latent reasoning approaches involve spending extra compute by doing "CoT in activation space," e.g. by grafting LM activations into other layers/later forward passes through the same model; our approach can be viewed as doing exactly the same thing, but instead "outsourcing" the CoT to another model (and thus reaping benefits from greater diversity of thoughts/reasoning paths from distinct models).

Also, note that as shown above, we find that even when models have different architectures, tokenization schemes, or training distributions, we do not need to train a mapping matrix $\mathbf{W}$ to attain SOTA results.

We hope that with the provided responses & additional strong results, you consider raising your score to support clear acceptance. Please let us know if there are any additional questions or concerns; we'd be happy to address them.

[1] Chau Pham, Boyi Liu, Yingxiang Yang, Zhengyu Chen, Tianyi Liu, Jianbo Yuan, Bryan A. Plummer, Zhaoran Wang, and Hongxia Yang. Let models speak ciphers: Multiagent debate through embeddings, 2024.

[2] Abhimanyu Dubey et al. The llama 3 herd of models, 2024.

[3] Vapnik, V. N. An overview of statistical learning theory, 1999.

We thank the reviewer for their feedback, and address their comments below:

1. The paper considers rather simple datasets (2 multiplayer and 7 reasoning tasks) and two models (llama 3b and 8b).

In Appendix B.6, we display results on the entire MMLU benchmark (57 datasets), spanning various domains and difficulty levels. We find that AC matches/outperforms NLD on 48/57 datasets, demonstrating our approach's value.

While our initial experiments were only with LLaMA models, note: (1) These LMs are among SOTA open-source models, meriting our focus; and (2) we extensively vary both the LLaMA suite (LLaMA-2, 3, 3.1, 3.2) and parameter count (1-70B) as shown in Figure 3 of the paper. This is already a broad coverage of models/model sizes.

However, we share additional results using models from the Qwen-2.5 and Gemma-2 families below. Each cell contains two results: Biographies score / GSM8k score.

Note the following:

• AC beats NLD across the board, and beats both individual models for 4/5 of the 6 model pairs on Biographies/GSM8k respectively----demonstrating the efficacy of AC irrespective of model architecture, size, tokenizer, and training data
• These results are obtained without training $\mathbf{W}$, meaning we do not need to train a projection layer between activation spaces to attain SOTA results, even for extremely distinct models! (We hypothesize this is because we are only replacing $B$'s last-token activation, hence $B$ can learn from $A$ without an extreme alteration to its activation distribution)

2. The experimental evaluation can be improved with better baselines than single model and NL debate.

We limit our focus to task-agnostic methods; existing model composition/grafting methods require substantial task-specific parameters/data & much more compute in the form of LM finetuning, layer or router training, etc., hence we do not compare to these.

Regarding multiagent debate, the NLD setup that we evaluate against is the predominant method of NL communication. This is quite a strong/flexible NL approach, involving CoT and allowing varied numbers of agents/rounds. (In fact, as shown in Appendix B.5, we find AC outperforms NLD even with 3-4 rounds of debate on 5 of the 7 reasoning datasets.)

3. Why was k = j = 26 selected?

See Section 3.3, lines 201-222. This reasoning is precisely why we choose a layer around halfway through the LM; indeed, Hernandez24 find that by around the halfway point of an LM's computation, it has developed "enriched entity representations" of the input that would be quite useful for communication compared to the next-token representations of later layers.

We validate this empirically; Figure 2 shows 2D contour plots of accuracy over different $k,j$ values. For computational reasons we only do this hyperparameter sweep on the Countries and Tip Sheets datasets and simply cross-apply the optimal values here, $k=j=26$, to the reasoning benchmarks and find that the values seem to generalize well across datasets, which is quite interesting in its own right.

5. The explanation in line 25-255 is not convincing about why the experimental comparison with Pham24 was avoided. Could you please elaborate further?

Please see point #4 in the reply to Reviewer gYwD below.

We hope that with the provided responses & additional strong results, you consider raising your score to support clear acceptance. We are happy to address any additional questions.