AlphaGo Moment for Model Architecture Discovery

ASI-Arch Main Results

📋 Overview

This is the official repository for our work "AlphaGo Moment for Model Architecture Discovery". We present a highly autonomous, multi-agent framework that empowers a Large Language Model (LLM) to conduct end-to-end scientific research in the challenging domain of linear attention mechanisms.

• 📖 Paper: Read our paper (Coming Soon)
• 💻 Code: Complete pipeline, database, and cognitive library available in this repository
• 🔬 Architectures: We open-source all 106 discovered state-of-the-art linear attention architectures

Our framework supports the following functionalities:

• Autonomous Architecture Discovery Pipeline: Complete multi-agent system that autonomously hypothesizes novel architectural concepts, implements them as code, and empirically validates their performance through systematic experimentation.

• Architecture Database: MongoDB-based cloud service storing all historical experimental data and enabling agent information retrieval, supporting multi-pipeline parallel execution.

• Cognition Base: MongoDB-powered knowledge repository providing relevant paper cognitions and research insights to guide the autonomous research process.

ASI-Arch autonomous research framework demonstrating AI's capability to conduct end-to-end scientific discovery, from hypothesis generation to empirical validation.

🏆 Performance

ASI-Arch has successfully discovered 106 novel linear attention architectures that achieve state-of-the-art performance across various benchmarks.

Performance comparison of 5 selected novel linear attention architectures discovered by ASI-Arch.

Our system demonstrates continuous optimization capability, consistently improving architecture quality throughout the autonomous research process.

Performance indicators showing steady improvement in benchmark scores and consistent reduction in loss values, with composite fitness scores demonstrating rapid initial improvement followed by gradual plateau.

📋 Table of Contents

• 🚀 Get Started
  • System Requirements
  • Installation
  • Environment Setup
• 🔧 Framework Components

🚀 Get Started

System Requirements

• Python 3.8+
• MongoDB 4.4+
• Docker & Docker Compose
• CUDA-compatible GPU (recommended)
• Minimum 16GB RAM, 32GB (recommended)
• WSL (recommended) with tmux (optional)

Installation

1. Clone the repository:

   git clone https://github.com/GAIR-NLP/ASI-Arch.git
   cd ASI-Arch

2. Install Dependencies:

   sudo apt update
   
   # Helpful tooling
   sudo apt install tmux -y
   sudo apt install nvitop
   curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash
   
   # Install NVIDIA CUDA Drivers - Example for Ubuntu 24.04
   sudo apt-get update
   sudo apt-get install -y wget gnupg
   wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
   sudo dpkg -i cuda-keyring_1.1-1_all.deb
   sudo apt-get update
   sudo apt-get -y install cuda
   
   # Pipeline Component (Main runner)
   conda create -n asi-arch python=3.10
   conda activate asi-arch
   pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu124
   pip install -r requirements.txt
   conda deactivate
   
   # Database Component
   conda create -n asi-arch-database python=3.10
   conda activate asi-arch-database
   pip install -r database/requirements.txt
   conda deactivate
   
   # Cognition Base Component
   conda create -n asi-arch-cognition-base python=3.10
   conda activate asi-arch-cognition-base
   pip install -r cognition_base/requirements.txt
   conda deactivate

Configuration Settings

Before running, you will want to first alter the configuration settings in config_agents.yaml and config_pipeline.yaml

Environment Setup

The framework relies on Docker to run the Database and Cognition Base services in the background. The steps below discuss how to do this in a single WSL window with 3 terminals open, using tmux

tmux
ctrl+b --> \" # Create a horizontal window
ctrl+b --> % # Split the bottom window vertically
ctrl+b --> % # Split the bottom window vertically (3rd window)

# Change the pane sizes to your liking
Ctrl+b  ↑   # expand the pane upward
Ctrl+b  ↓   # expand downward
Ctrl+b  ←   # expand left
Ctrl+b  →   # expand right

It will look something like this when set up effectively

# Window 2 (Bottom left):
# Start Cognition Base Service:
conda activate asi-arch-cognition-base
cd cognition_base
docker-compose up -d
python rag_api.py

# Window 3 (Bottom middle):
# Monitor GPU and CPU utilization
nvitop

# Window 4 (Bottom right):
# Start Database Service:
conda activate asi-arch-database
cd database
docker-compose up -d
./start_api.sh

# Window 1 (Top)
# Main Pipeline:
conda activate asi-arch
# Run the following commands on first-run or for fresh model discovery - more detailed instructions below
# python delete_all_database.py
# python seed_invismark.py
cd pipeline
python pipeline.py # Running Architecture Discovery

Seeding the initial environment

See Database Initialization Guide for more detailed dataabase seeding instructions.

🔧 Framework Components

ASI-Arch is composed of three interconnected systems working in concert to achieve autonomous scientific discovery.

🧬 Autonomous Architecture Discovery Pipeline (pipeline/)

The pipeline is the core engine of ASI-Arch, executing an autonomous loop of architectural innovation. It orchestrates a team of specialized agents to systematically hypothesize, implement, and validate new linear attention mechanisms.

• Core Modules:
  • evolve: The creative heart of the system. It generates novel architectural ideas by evolving existing designs. This module includes a Planner to design new models, a Code Checker to ensure correctness, and Deduplication agents to foster true innovation.
  • eval: The empirical validation module. It takes new architectures, trains them, and runs benchmarks. A Trainer agent handles the training process, while a Debugger agent can automatically analyze and fix errors during training.
  • analyse: The analysis module. An Analyzer agent provides a comprehensive breakdown of an experiment's results, comparing them to baselines and previous experiments to extract key insights.
• Workflow: The pipeline follows a continuous cycle: Sample an effective parent architecture from the database -> Evolve it into a new design -> Evaluate its performance through training and testing -> Analyze the results to generate insights -> Update the database with the new findings.
• Execution: The entire loop can be started with python pipeline/pipeline.py.

🗄️ Architecture Database (database/)

The database serves as the collective memory of the entire research process. It stores all historical experimental data, including architectures, results, analysis, and evolutionary lineage, enabling agents to learn from past successes and failures.

• Technology: Built on MongoDB for robust data storage, with a FastAPI server (mongodb_api.py) providing a comprehensive REST API for data access.
• Key Components:
  • mongodb_database.py: A high-level client for all database operations, managing the storage and retrieval of experimental DataElements.
  • candidate_manager.py: Maintains a curated list of the top-performing architectures (the "candidate set"). This elite set is used to guide the evolutionary process towards promising directions.
  • faiss_manager.py: Integrates FAISS for high-speed vector similarity search. This is crucial for the Deduplication agent to quickly check if a new idea is truly novel or just a rehash of a previous one.
  • evaluate_agent/: Contains a specialized Model Judger agent that provides a quantitative score for any given architecture based on its performance, innovation, and complexity.
• Usage: The database service is launched via database/start_api.sh, which starts the MongoDB container and the FastAPI server.

🧠 Cognition Base (cognition_base/)

The Cognition Base acts as the system's "domain expert," providing the agents with relevant knowledge from a vast library of scientific papers. It uses a Retrieval-Augmented Generation (RAG) approach to ground the agents' decisions in established research.

• Technology: Implemented as a RAG service (rag_service.py) that uses vector embeddings to find relevant passages from a corpus of research papers stored in the cognition/ directory.
• Key Components:
  • Knowledge Corpus: The cognition/ directory contains hundreds of JSON files, each representing a processed research paper on relevant topics like model architecture and attention mechanisms.
  • rag_service.py: The core service that loads the knowledge corpus, creates vector embeddings for the text, and uses a vector database (like OpenSearch) for efficient retrieval.
  • rag_api.py: A Flask-based API that allows other parts of the system, particularly the Analyzer and Planner agents in the pipeline, to query the knowledge base with natural language questions and receive the most relevant research insights.
• Usage: The RAG service is started via python cognition_base/rag_api.py, making the knowledge base available to the entire framework.