Communicating Activations Between Language Model Agents

Thank you for your valuable feedback! We appreciate that you found our approach original and well-motivated. We address specific comments from the review below:

The experiment set-up for the coordination game is not realistic—player A can simply send its text prompt to player B without incurring additional computation, since the models are identical. It would be more realistic if model A is smaller than model B in this case.

There is significant theoretical interest in this setup; both Countries and Tip Sheets are examples of both Lewis signaling games and two-player language coordination games which have been heavily studied in multi-agent communication literature [1, 2]. Also the core purpose of these toy setups is to show AC communication is indeed a viable method while being extremely computationally efficient. (In fact, if $A$ were to simply output its prompt and then pass that as context for $B$, this would require more compute than AC!)

Furthermore, in general (and as is shown by the reasoning benchmark setup), $B$ may need access to not just necessary information in $A$'s prompt but instead/also knowledge encoded in $A$'s weights.

An additional baseline for collaborative reasoning is CoT with model $B$ alone. This will illustrate whether the quality gains from A+B comes from additional inference computation.

Note that the single-model baselines in Table 3 already involve CoT prompting, showing that the quality gains from AC are independent of additional inference computation. We have made this clearer in the revised paper.

A useful direction for experiments is to investigate the case of RAG, where the retrieved contents are processed by a small model $A$ (whose activations are potentially pre-computed), while the response is generated by model $B$.

While this experiment would be too costly to run on our setup, this definitely makes sense, and is a great application of AC. We will plan to pursue this in future work.

Could we use cross attention to incorporate activations from model A to B?

This is exactly the method of CALM [3]. However in this work, we limit our focus to task-agnostic methods; CALM and other composition/grafting methods 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.

Does this suggest that the inputs to $f$ contain only the activations for the last tokens from $A$ and $B$ and the grafting only affects the activation of the last token of model B?

Yes. We have made this clearer in the revised paper.

The computation complexity analysis seems to be contradictory: "here $F(P, D) = O(PD)$ for non-learned $f$ and $O(D^2)$ when $f$ is the mapping matrix". $O(D^2)$ implies that only the last tokens are used, while $O(PD)$ seems to suggest that the activations of all prefix tokens are used.

Thank you, this should say "$\mathcal{F}(D) = O(D)$ for non-learned $f$ and $O(D^2)$ when $f$ is the mapping matrix". We have fixed this in the revised paper.

We hope these improvements adequately address your questions, and we hope that you consider the extent of these improvements in your evaluation. Please let us know if there are any additional questions or concerns; we'd be happy to address them.

[1] David Lewis. Convention: A philosophical study, 2008.

[2] Angeliki Lazaridou, Alexander Peysakhovich, and Marco Baroni. Multi-agent cooperation and the emergence of (natural) language, 2016.

[3] Rachit Bansal, Bidisha Samanta, Siddharth Dalmia, Nitish Gupta, Shikhar Vashishth, Sriram Ganapathy, Abhishek Bapna, Prateek Jain, and Partha Talukdar. Llm augmented llms: Expanding capabilities through composition, 2024.

Thank you for your valuable feedback! We appreciate that you found our approach interesting and well-motivated, and our results and analysis useful. We address specific comments from the review below:

The approach only works for different model weights. While it is true that the paradigm of debate between different models is interesting, it still is useful to support debate between identical models.

Thanks for raising this point. We can extend AC to work for debate between identical models as follows: let $A$ and $B$ be instances of the same model. We can sample a completion from $A$ with temperature and graft the last-token layer-$k$ activation of the completion into $B$ at layer $j$ as part of the AC procedure. This still saves a substantial amount of compute over NLD between 2 model instances, showing our technique can apply to this setting. Below are the results of this experiment ($512$-token completions sampled with temperature $0.7$, debate run for $2$ rounds with $2$ instances of $`LLaMA-3.1-8B`$, and AC run with $A = B = `LLaMA-3.1-8B`$):

Indeed, while communication between multiple model instances doesn't always show improvement over the single model itself (a well-known result from [1]), AC matches/outperforms NLD on five of the seven datasets.

The intuition behind debate between multiple identical model instances is that sampling multiple completions (with temperature) from the same model yields diverse reasoning paths that can be recombined into a stronger final answer. The above experiment shows that the same intuition holds for AC—we are sampling multiple times from the same model, but passing responses between agents via AC rather than as NL messages.

More broadly, however, we claim that the multiple-model setting is just as 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 and of different sizes 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?

Thank you for your valuable feedback! We appreciate that you found our approach effective and compute analysis strong. We address specific comments from the review below:

The last token activation replacement cannot capture the nuanced response from CoT traces. This might explain the GSM8k performance where the proposed method does not perform well compared to NLD baseline since it requires intermediate steps for answer compuation.

We distinguish between two questions here, and address both in turn:

Is replacing just the last token activation enough to communicate information from $A$ to $B$?

Note that 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. Indeed, this must be the case for activation communication to recover the gap between the zero-communication and skyline setups on both coordination games, which (for Tip Sheets in particular) require information starting at the first few tokens of $A$'s prompt to be communicated.

In addition, we experiment with summing the activations of all tokens in the sequence rather than just the last (we cannot replace all tokens as this would just replace $B$'s layer-$j$ activation with $A$'s layer $k$-activation):

Indeed, applying $f$ to all tokens decreases performance relative to applying $f$ to just the last token. Note that the fact performance generally decreases from $f =$ replace to $f =$ sum, and further with all tokens, is expected. The high performance of AC with $f =$ replace means that the edited last-token activation $\mathbf{b}$ retains some meaning in $B$'s activation space; it is less likely for this to be the case when $f =$ sum (at the very least $\mathbf{b}$ has norm roughly $2\times$ that of $B$'s original last-token activation), and when doing this for all tokens we'd expect performance to decrease even further as now all activation vectors, not just the last, are out-of-distribution w.r.t. $B$'s activation space.

Does activation communication lose out on the benefits of CoT? [This is also the reviewer's first question]

This is a great question, and we empirically verify the answer is no.

First, note that we still reap the benefits of CoT when we sample a completion from $B$ after AC (where $B$ gets all the information encoding $A$'s "beliefs" about the prompt via AC, hence CoT on $A$'s side is not needed). To verify this intuition, we experiment with prompting $A$ with CoT, generating a full response, and then passing the layer-$k$ last-token activation of the CoT response to $B$ as part of AC:

Indeed, we empirically find our above intuition to hold, as there is no significant improvement over vanilla AC when generating from $A$ using CoT.

Thank you for your valuable feedback! We appreciate that you found our approach original and well-supported and our experiments thorough. We address specific comments from the review below:

Interpretability and Debugging Concerns: While we acknowledge the complexity of "debugging AC", we would like to make a few points. First, past work has justified that we do not necessarily need to understand the intermediate debates amongst LLMs [1]. Second, there is a fundamental tradeoff between interpretability and information preservation (as activations, by virtue of being much higher-dimensional than the space of natural language, allow proportionally higher-entropy communication) [1], which merits discussion beyond the scope of this work. But third, we actually posit that our method suggests a new avenue towards interpreting LM activations: "translating" activations based on the beliefs they induce as messages in listening agents, similar to the method put forward in [2]. This would make the communicated activation vectors interpretable, and we recognize this as a promising avenue for future research.

Sparse Performance Gains on Some Tasks: We'd like to push back against this claim, as the consistency of AC outperforming single-model baselines and NLD is well-supported by our results in Table 3. To reiterate our analysis from the paper: Notably, while NLD consistently outperforms $`LLaMA-3.2-3B`$, it does not always display a performance improvement over $`LLaMA-3.1-8B`$; but remarkably, AC consistently outperforms both single-model baselines. In fact, AC offers an up to $27.0$% improvement over NLD across six of the seven reasoning datasets.

Selection of Baselines: We limit our focus to task-agnostic methods; approaches such as CALM [3] and other composition/grafting methods 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.

Regarding multiagent debate, the NLD setup from [4] that we evaluate against is the predominant method of NL communication. This is quite a strong and flexible NL approach, incorporating CoT into each model's prompt and allowing varied numbers of agents and rounds. Indeed, we find in additional experiments that AC outperforms NLD even with additional rounds of debate:

We see that 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.

Scalability and Real-world Applicability: Exploring API-only approaches is highly limiting (see AC's meta review of [1] at last year's conference). Furthermore, recent releases of powerful open-source models [5] merit the development of embedding-based techniques.

As shown above, AC matches/outperforms NLD on 48/57 datasets, indicating its superiority and robustness as an alternative "language" for inter-LLM communication.

We hope that with these strong and complete benchmark results now provided, the reviewers consider raising their scores to support acceptance.