MemCollab: Collaborative Memory Framework

• MemCollab is a collaborative memory framework that shares agent-invariant reasoning constraints across heterogeneous LLM-based agents using contrastive distillation.
• It builds a shared, agent-agnostic memory that enhances accuracy and efficiency in tasks like mathematical reasoning and code generation.
• The framework employs a task-aware retrieval mechanism that filters out agent-specific biases and reduces reasoning steps during inference.

MemCollab is a collaborative memory framework designed to facilitate cross-agent knowledge sharing among heterogeneous LLM-based agents. Unlike conventional agent-specific memory systems, MemCollab constructs a shared, agent-agnostic memory by applying contrastive distillation to reasoning trajectories produced by diverse agents over the same tasks. This shared memory encodes abstract reasoning constraints—specifically, task-level invariants that generalize across different models—while filtering out agent-specific idiosyncrasies. At inference, a task-aware retrieval mechanism selects and integrates the most relevant memory entries to condition agent reasoning, resulting in improved accuracy and efficiency across mathematical reasoning and code generation tasks. Empirical evaluation demonstrates consistent gains in performance and inference efficiency, including in cross-model-family deployments, thereby substantiating the utility of MemCollab as a universal reasoning resource for LLM-based agents (Chang et al., 24 Mar 2026).

1. Components and Architecture

MemCollab comprises three principal components: (1) a heterogeneous agent pool $\mathcal{A} = \{A_i\}_{i=1}^N$, where each agent $A_i$ corresponds to a distinct LLM and reasoning paradigm; (2) a shared memory bank $\mathcal{M}$, populated with distilled agent-agnostic constraints derived via contrastive analysis; and (3) a task-aware retrieval interface that filters and ranks memory entries contingent on the current query's task category.

The high-level workflow entails, during training: (1) independent generation of reasoning trajectories $\tau_i^{(x)}$ by all agents on each task $x$; (2) contrastive distillation of shared invariants from these trajectories, forming memory entries $\{m_k^{(x)}\}$ aggregated into $\mathcal{M}$. At test time, given a query $q$, the retrieval mechanism classifies $q$ into a task category, retrieves the most relevant entries $\mathcal{M}_q \subset \mathcal{M}$, and conditions the agent's solution on $\mathcal{M}_q$.

┌─────────────────┐ │ Training Set │ │ tasks x∈𝒟 │ └─────────────────┘ │ ┌──────▼────────┐ │Multiple Agents│ │ A₁,…,A_N solve│ │ x → τ₁, …, τ_N│ └──────┬────────┘ │ ┌──────▼──────────────┐ ┌─────────────────────┐ │Contrastive │───▶│ Memory Bank │ │Trajectory │ │ ℳ ← {m_k^(x)} │ │Distillation │ └─────────────────────┘ └─────────────────────┘ ▲ ▲ ┌───────────────┐ ┌────────────┴────────────┐ │ Inference │ │ Retrieval │ │ Task q │<───┤ Interface │ │ A(q, ℳ_q) │ └───────────────────────┘ └───────────────┘

2. Formalization and Objective

Let $\mathcal{A} = \{A_i\}_{i=1}^N$ be the set of agents and $\mathcal{D} = \{x\}$ the training tasks. Each agent $A_i$ produces a reasoning trajectory on task $x$ given by $\tau_i^{(x)} = (r_{i,1}, \ldots, r_{i,T_i})$. Memory entries $m_k$ encode pairs $(\text{invariant}; \text{forbidden pattern})$, capturing essential constraints and anti-patterns.

The principal objective is to build a shared memory $\mathcal{M}$ that:

1. Encapsulates task-relevant, agent-invariant knowledge;
2. Suppresses agent-specific biases;
3. Remains usable by any $A_i$ to enhance both accuracy and inference-time efficiency.

3. Contrastive Trajectory Distillation

Each trajectory $\tau$ is embedded via an agent-specific encoder $f_i: \tau_i \mapsto h_i \in \mathbb{R}^d$, composed with a projection head $g:\mathbb{R}^d \to \mathbb{R}^{d'}$ to yield $z^t_i = g(f_i(\tau^t_i))$ per trajectory step $t$.

Contrastive learning is applied: positive pairs $\{z^t_i, z^t_j\}$ for agents $i \neq j$ on matching steps, negatives from unrelated or erroneous trajectories. The InfoNCE loss is: $$L_{\rm ctr} = -\sum_{t=1}^T \log \frac{\exp\!\bigl(\mathrm{sim}(z^t_i,z^t_j)/\tau\bigr)} {\sum_{k=1}^N \exp\!\bigl(\mathrm{sim}(z^t_i,z^t_k)/\tau\bigr)}$$ where $\mathrm{sim}(u,v)$ denotes cosine similarity and $\tau$ is the temperature. Minimization of $L_{\rm ctr}$ aligns agent-invariant features, yielding distilled invariants for $\mathcal{M}$.

4. Memory Bank Construction and Indexing

Memory is stored as a key-value set: $$\mathcal{M} = \{ (k_j, v_j) \}_{j=1}^{|\mathcal{M}|}$$ where $k_j \in \mathbb{R}^d$ is an embedding derived from invariant summarization, and $v_j$ is a natural language specification (e.g., “enforce <invariant>; avoid <pattern>”).

Each entry is further labeled by task category $(c_j, u_j)$ for main and subcategory identification. An inverted index enables efficient retrieval by $(c_j, u_j)$, and within each category, $k_j$ are indexed for approximate nearest-neighbor search (e.g., FAISS).

5. Task-Aware Memory Retrieval

A classifier $h: x \mapsto (c_x, u_x)$ assigns each query to its task category and subcategory. Retrieval uses the key $k_q = f_q(q)$ and scores each memory record $(k_j, v_j, c_j, u_j)$ via: $$\mathrm{score}(q, k_j) = \mathrm{sim}(k_q, k_j) + \lambda \mathbf{1}[c_q = c_j \wedge u_q = u_j]$$ Retrieval selects the top-$K$ scoring entries, which are prepended or interleaved into the agent’s input (prompt), thus strongly biasing the agent's decoding toward enforcing distilled invariants while avoiding forbidden patterns.

6. Training and Inference Protocol

Memory Construction

Pseudocode:

Input: Agents {A_i}, training tasks 𝒟, max entries K Output: Memory bank ℳ ℳ ← ∅ for x in 𝒟 do for each i≠j in Agents do τ_i ← A_i(x) τ_j ← A_j(x) # Encode or summarize trajectories z_i ← f_i(τ_i); z_j ← f_j(τ_j) # Contrastive update update_encoders_via_Lctr(z_i,z_j) end # After convergence, extract up to K invariants {m₁,…,m_K} ← summarize_invariants({τ_i}_{i}, {z_i}_{i}) add {m₁,…,m_K} to ℳ end return ℳ

Inference and Retrieval

Pseudocode:

Input: query q, memory ℳ, retrieval count p, agent A Output: answer y (c_q,u_q) ← h(q) # classify q’s task ℳ_q ← Top−p({ (k_j,v_j)∈ℳ : (c_j,u_j)=(c_q,u_q) }) y ← A.solve(q; prepended=ℳ_q) return y

7. Empirical Evaluation and Observed Performance

Experiments were conducted on mathematical reasoning (MATH500, GSM8K) and code generation (MBPP, HumanEval) benchmarks. Key performance metrics included accuracy / pass@1 and mean number of reasoning turns.

Key results:

• The Qwen-2.5-7B model's accuracy increased from 52.2% to 67.0% on MATH500 and from 47.9% to 57.6% on MBPP, outperforming both baseline and naïve cross-agent memory transfer approaches.
• Larger models (Qwen-2.5-32B) achieved absolute improvements of +3–5% on all tasks.
• Memory distilled from cross-architecture (LLaMA3-8B ↔ Qwen-32B) agent pairs yielded comparable or superior results compared to within-family pairs, confirming effective agent-agnostic constraint extraction.
• Average reasoning turns decreased (e.g., from 2.7 to 2.2 on MATH500), evidencing improved inference efficiency.

These results substantiate MemCollab's efficacy in constructing collaborative, agent-invariant memory that enhances both the problem-solving capability and efficiency of heterogeneous agent ensembles (Chang et al., 24 Mar 2026).