The Problem

Multi‑agent reasoning often explodes token usage because each agent writes full textual traces. This makes inference slow and costly, especially when large language models are used as back‑ends.

What This Does

LatentMAS moves the collaboration into the model’s latent space. Agents keep a working‑memory tensor and exchange “latent thoughts” instead of strings. The core implementation lives in:

  • methods/latent_mas.py – latent‑space MAS logic.
  • methods/text_mas.py – text‑based baseline for comparison.
  • methods/baseline.py – simple non‑latent baseline.
  • models.py – model wrappers, token‑level utilities, and rendering.
  • utils.py – data‑loading, seeding, and answer normalisation helpers.

The repository ships with data loaders (data.py), prompt builders (prompts.py), and a small CLI driver (run.py).

How It Is Wired

Execution starts at run.py → main (line 84). main parses a command‑line batch (the parser is not shown but process_batch is the next entry point) and then calls process_batch, which:

  1. Loads data via load_gsm8k, load_aime2025, etc. in data.py.
  2. Builds agent messages using prompts.build_agent_message_* functions.
  3. Selects a method (default_agents from methods/__init__.py returns a list of Agent objects).
  4. Runs the batch by invoking run_batch (or run_batch_vllm) in the selected method class (methods/latent_mas.py or methods/text_mas.py).

run_batch is the most connected function (6 outgoing imports, instability 1). It repeatedly calls:

  • normalize_answer (10 call sites) – central to evaluation.
  • _past_length (5 call sites) – computes latent memory size.
  • prepare_chat_batch, generate_text_batch, run_with_timeout – handle token‑level interaction and safety timeouts.

run_batch also calls extract_gsm8k_answer (used in 2 places) and render_chat (in models.py).

The utility hub is utils.py (instability 0, Ca 5). It provides set_seed, auto_device, and answer‑extraction helpers used across all methods.

models.py supplies the model wrapper (_ensure_pad_token, _past_length, render_chat, prepare_chat_input) and is called from three other files.

Because run reaches 42 functions and process_batch reaches 23, the call graph is relatively shallow: the bulk of work stays inside run_batch → helper utilities → model wrapper. No circular imports are present, simplifying future refactors.

How To Use It

# Clone the repo
git clone https://github.com/moses-y/LatentMAS.git
cd LatentMAS

# Install Python dependencies
pip install -r requirements.txt

# Run a benchmark batch (example uses the default data loader)
python run.py

No additional configuration files, environment variables, or secret keys are required. The CLI reads the batch definition from the script itself (the README does not specify external config).

Real‑World Use

A service that needs fast multi‑step reasoning (e.g., code‑generation assistants) can replace a token‑heavy chain of LLM calls with a single latent‑memory pass:

from run import process_batch
process_batch()   # loads GSM8K, runs latent agents, returns answers

The returned answers are already normalised by utils.normalize_answer, ready for downstream scoring or user display.

Code Health & Issues

  • HIGH – Duplicated code – identical 6‑line blocks appear in methods/baseline.py, methods/latent_mas.py, methods/text_mas.py. Consolidate into shared helpers.
  • MEDIUM – Broad exception handling – utils.py catches generic except: statements, risking silent failures. Replace with specific exception types and proper logging.
  • MEDIUM – Deep nesting – nesting depth of 7 in methods/latent_mas.py, methods/text_mas.py, and run.py makes the control flow hard to follow; refactor with early returns or helper extraction.
  • MEDIUM – No test suite – repository contains no tests/ directory; untested paths increase regression risk.
  • MEDIUM – No CI/CD – absence of .github/ workflows or other automation means builds are not automatically validated.
  • LOW – No lockfile – requirements.txt is present but no requirements.lock; reproducible builds depend on external version resolution.

The Bottom Line

LatentMAS delivers a functional prototype for latent‑space multi‑agent reasoning with a clear entry point and modest dependency footprint. The code base is small enough to understand quickly, but duplicated logic, broad exception catches, and deep nesting increase maintenance overhead. Teams that need a research‑grade baseline can adopt it now, but should add tests, CI, and refactor shared code before extending it in production.