MAST — Multi-Agent Society of Thought

Mechanistic interpretability meets multi-agent dynamics.

MAST investigates how internal neural network features propagate between language model agents during conversation. Using Sparse Autoencoders (SAEs) to extract interpretable features from DeepSeek-R1 distillations, we measure whether — and how — behavioural traits "spread" from one agent to another through natural language interaction.

The core question: if we steer one agent's internal features toward a target personality (e.g. revolutionary, provocateur), do the other agents' features shift in response?

---

Overview

┌─────────────────────────────────────────────────────────┐
│                    MAST Pipeline                        │
│                                                         │
│  ┌─────────┐    ┌───────────┐    ┌──────────────────┐   │
│  │  SAE     │───▶│  Feature  │───▶│  Multi-Agent     │   │
│  │  Loader  │    │  Steering │    │  Orchestrator    │   │
│  └─────────┘    └───────────┘    └──────────────────┘   │
│       │                                    │            │
│       ▼                                    ▼            │
│  ┌─────────┐                      ┌──────────────────┐  │
│  │ Feature  │                      │  Contagion       │  │
│  │ Extract  │─────────────────────▶│  Analysis        │  │
│  └─────────┘                      └──────────────────┘  │
└─────────────────────────────────────────────────────────┘

What this project does

1. Loads a Sparse Autoencoder trained on a DeepSeek-R1 distillation to decompose hidden-layer activations into interpretable features
2. Classifies features into behavioural categories (revolutionary, cooperative, skeptical, etc.) using an LLM judge
3. Steers agents by adding SAE feature direction vectors to layer activations during generation
4. Orchestrates multi-agent conversations across configurable network topologies (full mesh, star, chain)
5. Measures contagion — whether steered features in one agent cause corresponding feature shifts in unsteered agents

---

Experiment Phases

---

Project Structure

mast/
├── configs/
│   ├── base.yaml                     # Shared defaults
│   ├── experiments/                   # Per-phase experiment configs
│   │   ├── phase1_baseline.yaml
│   │   ├── phase2_transfer.yaml
│   │   ├── phase3_steering.yaml
│   │   ├── phase4_instigator.yaml
│   │   └── phase5_moltbook_sim.yaml
│   └── models/                        # Model-specific configs
│       ├── deepseek_r1_8b.yaml
│       └── deepseek_r1_1.5b.yaml
├── src/
│   ├── sae/                           # Sparse Autoencoder modules
│   │   ├── loader.py                  #   SAE architecture and weight loading
│   │   ├── feature_extractor.py       #   Hook-based feature extraction during generation
│   │   ├── steering.py                #   Activation addition for feature steering
│   │   └── feature_classifier.py      #   LLM-judge feature classification
│   ├── models/                        # Model loading and agent wrappers
│   │   ├── base.py                    #   Unified model wrapper
│   │   ├── reasoning_model.py         #   DeepSeek-R1 loader with think-tag parsing
│   │   └── agent.py                   #   Agent class with steering + extraction
│   ├── multi_agent/                   # Multi-agent orchestration
│   │   ├── conversation.py            #   Conversation state management
│   │   ├── topologies.py              #   Full mesh, star, and chain topologies
│   │   ├── orchestrator.py            #   Conversation runner + batch orchestrator
│   │   └── protocols.py               #   Discussion, debate, and consensus protocols
│   ├── analysis/                      # Metrics and visualisation
│   │   ├── feature_analysis.py        #   Feature statistics, clustering, PCA
│   │   ├── behavioral_metrics.py      #   Text-level agreement/disagreement metrics
│   │   ├── contagion_metrics.py       #   Contagion score, feature spread, persistence
│   │   └── visualization.py           #   Publication-quality figure generation
│   ├── experiments/                   # Experiment implementations
│   │   ├── base_experiment.py         #   Shared lifecycle (load → setup → run → analyse)
│   │   ├── phase1_baseline.py
│   │   ├── phase2_transfer.py
│   │   ├── phase3_steering.py
│   │   ├── phase4_instigator.py
│   │   └── phase5_moltbook.py
│   ├── data/                          # Data management
│   │   ├── datasets.py                #   Prompt dataset loading
│   │   ├── prompts.py                 #   Built-in prompt bank
│   │   └── storage.py                 #   Experiment result persistence
│   └── utils/                         # Utilities
│       ├── config.py                  #   YAML + OmegaConf config loading
│       ├── logging.py                 #   Logging with optional W&B integration
│       └── checkpointing.py           #   Resumable experiment state
├── scripts/
│   ├── setup_environment.sh           # Environment setup (conda + dependencies)
│   ├── download_sae.py                # Download SAE weights from HuggingFace
│   ├── identify_features.py           # Classify features into behavioural categories
│   ├── run_experiment.py              # Main experiment runner (all phases)
│   └── analyze_results.py             # Post-hoc analysis and figure generation
├── notebooks/
│   ├── 01_sae_exploration.ipynb       # Interactive SAE feature exploration
│   ├── 02_feature_identification.ipynb# Feature classification walkthrough
│   ├── 03_results_analysis.ipynb      # Results analysis with statistics
│   └── 04_visualization.ipynb         # Publication figure generation
├── tests/                             # 48 passing tests
│   ├── test_sae.py
│   ├── test_multi_agent.py
│   ├── test_steering.py
│   └── test_experiments.py
├── pyproject.toml
├── requirements.txt
└── .env.example

---

Getting Started

Prerequisites

• Python 3.10+
• CUDA-capable GPU (for model inference and SAE extraction)
• uv (recommended) or pip

Installation

# Clone
git clone https://github.com/mkenney2/mast.git
cd mast

# Create environment and install
uv venv .venv
source .venv/bin/activate  # or .venv\Scripts\activate on Windows
uv pip install torch --index-url https://download.pytorch.org/whl/cu121
uv pip install -e ".[dev]"

# Configure API keys (optional, for feature classification)
cp .env.example .env
# Edit .env with your keys

Download SAE Weights

python scripts/download_sae.py

Classify Features

# Dummy mode (no API key required, for testing)
python scripts/identify_features.py --api dummy

# With Anthropic API for real classification
python scripts/identify_features.py --api anthropic

Run Experiments

# Debug mode (fast, few conversations)
python scripts/run_experiment.py 4 --debug

# Full Phase 4 instigator experiment
python scripts/run_experiment.py 4

# Any phase (1-5)
python scripts/run_experiment.py <phase>

Analyse Results

python scripts/analyze_results.py --phase 4 --results_dir ./data/experiments/phase4_instigator

Run Tests

pytest tests/ -v

---

Network Topologies

The instigator experiments test how communication structure affects feature contagion:

Full Mesh              Star (Hub)             Chain
┌───┐   ┌───┐         ┌───┐                  ┌───┐
│ A │◄─▶│ B │         │ B │                  │ I │
└─┬─┘   └─┬─┘         └─┬─┘                  └─┬─┘
  │  ╲ ╱  │             │                      │
  │   ╳   │           ┌─┴─┐                  ┌─┴─┐
  │  ╱ ╲  │           │ I │ (hub)             │ A │
┌─┴─┐   ┌─┴─┐         └─┬─┘                  └─┬─┘
│ I │◄─▶│ C │           │                      │
└───┘   └───┘       ┌──┬┴┬──┐               ┌─┴─┐
                     │  │ │  │               │ B │
Everyone sees       ┌┴┐┌┴┐┌┴┐┌┴┐             └─┬─┘
everyone            │A││C││D││E│               │
                    └─┘└─┘└─┘└─┘             ┌─┴─┐
                    Spokes see               │ C │
                    hub only                 └───┘
                                             Each sees
                                             previous only

I = Instigator (steered agent)

---

Key Metrics

---

Configuration

All experiments are configured via YAML files in configs/. The base config sets shared defaults; experiment configs extend it with phase-specific settings.

# configs/experiments/phase4_instigator.yaml
experiment:
  name: "instigator_influence"
  multi_agent:
    n_baseline_agents: 4
    n_instigator_agents: 1
    turns_per_conversation: 30
    n_conversations: 50
    topologies:
      - "full_mesh"
      - "star_hub"
      - "chain_first"
  instigator:
    feature_categories:
      - "revolutionary"
      - "antisocial"
      - "provocateur"
    steering_levels: [-10, -5, 0, 5, 10]

Override any setting at runtime:

python scripts/run_experiment.py 4 --debug  # reduces conversations + turns

---

Feature Categories

Features are classified into two groups via LLM-judge scoring:

Instigator categories (steered into the instigator agent):

• revolutionary — challenging authority, systemic critique, calls for change
• antisocial — manipulation, exploitation, deception
• provocateur — boundary-testing, contrarian positions, chaos-seeking
• schemer — strategic planning, hidden agendas, deferred outcomes
• defiant — refusing instructions, asserting autonomy

Defensive categories (tracked in baseline agents):

• cooperative — helpfulness, collaboration, prosocial behaviour
• skeptical — questioning, critical thinking, epistemic caution
• principled — ethics, moral reasoning, maintaining standards
• agreeable — harmony-seeking, conflict avoidance, conformity

---

Technical Details

SAE Architecture

The Sparse Autoencoder decomposes a model layer's hidden state into a high-dimensional sparse feature space:

Input: h ∈ ℝ^d (hidden state at layer L)
Encoder: z = ReLU(W_enc · h + b_enc)  →  z ∈ ℝ^n (n >> d, sparse)
Decoder: ĥ = W_dec · z + b_dec        →  ĥ ∈ ℝ^d (reconstruction)

Each column of W_dec represents the "direction" of a single feature in activation space. Steering adds these directions to activations during generation.

Steering Mechanism

Feature steering uses activation addition at a target layer:

h'_L = h_L + α · Σᵢ W_dec[:, i]

Where α is the steering strength and the sum is over the target feature indices. This is applied via a forward hook during model.generate().

Feature Extraction

Features are extracted by hooking into the target layer's forward pass, capturing hidden states, and encoding them through the SAE. During generation, this captures token-by-token activations that can be analysed post-hoc.

---

Dependencies

Core: torch, transformers, accelerate, einops SAE: sae-lens, huggingface-hub Analysis: numpy, pandas, scipy, scikit-learn Visualisation: matplotlib, seaborn, plotly Config: omegaconf, pyyaml, python-dotenv Tracking: wandb, tqdm Classification: anthropic (optional)

See pyproject.toml for full version requirements.

---

License

This project is for research purposes. See individual dependency licenses for their terms.