代码地址: https://github.com/GAIR-NLP/ASI-Arch?tab=readme-ov-file
1. 项目概述
ASI-Arch是一个高度自主的多智能体框架,旨在实现模型架构发现的"AlphaGo时刻"。该框架使大型语言模型(LLM)能够在具有挑战性的线性注意力机制领域进行端到端的科学研究。项目的核心目标是通过自主假设生成、代码实现和系统性实验验证来发现新颖的神经网络架构。
从整体架构来看,ASI-Arch由三个相互关联的核心系统组成:
- 自主架构发现管道(pipeline/)
- 架构数据库(database/)
- 认知基础(cognition_base/)
这三个系统协同工作,形成一个完整的闭环,使AI能够进行自主科学发现。
2. 核心模块详细分析
2.1 自主架构发现管道(pipeline/)
管道是ASI-Arch的核心引擎,执行自主架构创新的循环。它协调一组专门的智能体来系统地假设、实现和验证新的线性注意力机制。
2.1.1 核心子模块
-
evolve模块:系统的创新核心,负责通过进化现有设计来生成新颖的架构理念。
Planner
:设计新模型Code Checker
:确保代码正确性Deduplication
:避免重复创新,促进真正的创新
-
eval模块:实证验证模块,负责训练新架构并运行基准测试。
Trainer
:处理训练过程Debugger
:自动分析和修复训练中的错误
-
analyse模块:分析模块,
Analyzer
智能体提供实验结果的全面分析,与基线和先前实验进行比较以提取关键见解。
2.1.2 工作流程
管道遵循连续循环:
- 采样:从数据库中选择有效的父架构
- 进化:将其进化为新设计
- 评估:通过训练和测试验证其性能
- 分析:分析结果以生成见解
- 更新:将新发现更新到数据库
2.1.3 核心实现
从pipeline.py
中可以看到,主循环由main()
函数控制,持续执行实验:
async def main():
"""Main function - continuous experiment execution."""
set_tracing_disabled(True)
log_info("Starting continuous experiment pipeline...")
# Run plot.py first
log_info("Running plot scripts...")
log_info("Plot scripts completed")
experiment_count = 0
while True:
try:
experiment_count += 1
log_info(f"Starting experiment {experiment_count}")
success = await run_single_experiment()
if success:
log_info(f"Experiment {experiment_count} completed successfully, starting next experiment...")
else:
log_warning(f"Experiment {experiment_count} failed, retrying in 60 seconds...")
await asyncio.sleep(60)
except KeyboardInterrupt:
log_warning("Continuous experiment interrupted by user")
break
except Exception as e:
log_error(f"Main loop unexpected error: {e}")
log_info("Retrying in 60 seconds...")
await asyncio.sleep(60)
每个实验的具体执行由run_single_experiment()
函数处理,该函数按顺序调用各个模块:
async def run_single_experiment() -> bool:
# Start a new pipeline process
pipeline_id = start_pipeline("experiment")
try:
# Step 1: Program sampling
log_step("Program Sampling", "Start sampling program from database")
context, parent = await program_sample()
# Step 2: Evolution
log_step("Program Evolution", "Start evolving new program")
name, motivation = await evolve(context)
# Step 3: Evaluation
log_step("Program Evaluation", f"Start evaluating program {name}")
success = await evaluation(name, motivation)
# Step 4: Analysis
log_step("Result Analysis", f"Start analyzing program {name} results")
result = await analyse(name, motivation, parent=parent)
# Step 5: Update database
log_step("Database Update", "Update results to database")
update(result)
# Successfully complete pipeline
log_info("Experiment pipeline completed successfully")
end_pipeline(True, f"Experiment completed successfully, program: {name}, result: {result}")
return True
# 异常处理代码省略
2.2 架构数据库(database/)
数据库作为整个研究过程的集体记忆,存储所有历史实验数据,包括架构、结果、分析和进化谱系,使智能体能够从过去的成功和失败中学习。
2.2.1 技术栈
- 基于MongoDB进行稳健的数据存储
- FastAPI服务器(
mongodb_api.py
)提供全面的REST API用于数据访问
2.2.2 关键组件
mongodb_database.py
:所有数据库操作的高级客户端,管理实验DataElement
的存储和检索candidate_manager.py
:维护表现最佳架构的精选列表(“候选集”),用于引导进化过程faiss_manager.py
:集成FAISS进行高速向量相似性搜索,对Deduplication
智能体至关重要evaluate_agent/
:包含专门的Model Judger
智能体,基于性能、创新性和复杂性为任何给定架构提供量化评分
2.2.3 使用方式
通过database/start_api.sh
启动数据库服务,该脚本启动MongoDB容器和FastAPI服务器。
2.3 认知基础(cognition_base/)
认知基础充当系统的"领域专家",为智能体提供来自大量科学论文库的相关知识。它使用检索增强生成(RAG)方法将智能体的决策建立在已建立的研究基础上。
2.3.1 技术实现
- 实现为RAG服务(
rag_service.py
),使用向量嵌入从存储在cognition/
目录中的研究论文语料库中查找相关段落 - Flask-based API(
rag_api.py
)允许系统其他部分查询知识库
2.3.2 关键组件
- 知识语料库:
cognition/
目录包含数百个JSON文件,每个文件代表一篇关于模型架构和注意力机制等相关主题的处理过的研究论文 rag_service.py
:核心服务,加载知识语料库,为文本创建向量嵌入,并使用向量数据库(如OpenSearch)进行高效检索rag_api.py
:允许管道中的Analyzer
和Planner
智能体以自然语言查询知识库并接收最相关的研究见解
2.3.3 使用方式
通过python cognition_base/rag_api.py
启动RAG服务,使知识库可供整个框架使用。
3. 技术栈与实现细节
3.1 主要技术和框架
- 编程语言:Python 3.8+
- 数据库:MongoDB 4.4+,用于数据存储;FAISS用于向量相似性搜索
- 容器化:Docker & Docker Compose
- 深度学习框架:PyTorch 2.4.0
- API开发:FastAPI(数据库),Flask(认知基础)
- 异步编程:asyncio,用于处理异步操作
- 向量数据库:OpenSearch(认知基础)
3.2 代码组织方式
项目采用模块化设计,每个核心系统(管道、数据库、认知基础)都有自己的目录结构和明确的职责划分。代码遵循Python最佳实践,使用类型提示增强可读性和可维护性。
3.3 关键设计模式
- 智能体模式:系统中的各个组件(Planner, Analyzer, Trainer等)都被设计为专门的智能体,具有特定的职责和能力
- 管道模式:整个架构发现过程被组织为一个连续的管道,每个阶段都有明确的输入和输出
- 知识库模式:使用RAG技术构建的认知基础提供领域知识支持
- 数据库模式:使用MongoDB存储实验数据,支持复杂查询和分析
4. 运行示例
4.1 系统启动步骤
-
克隆仓库:
git clone https://github.com/GAIR-NLP/ASI-Arch.git cd ASI-Arch
-
创建Conda环境:
conda create -n asi-arch python=3.10 conda activate asi-arch
-
安装依赖:
# 安装主要依赖 pip install -r requirements.txt # 安装带CUDA支持的PyTorch pip3 install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu124 # 安装组件特定依赖 pip install -r database/requirements.txt pip install -r cognition_base/requirements.txt
-
启动数据库服务:
cd database docker-compose up -d ./start_api.sh
-
启动认知基础服务:
cd cognition_base docker-compose up -d python rag_api.py
-
运行架构发现:
conda activate asi-arch cd pipeline python pipeline.py
4.2 典型使用场景
- 自动架构发现:运行主管道自动发现新的线性注意力架构
- 性能评估:评估现有架构的性能并与基线比较
- 知识检索:查询认知基础获取相关研究论文的见解
- 数据分析:分析实验结果并提取关键见解
5. 性能与成果
ASI-Arch成功发现了106种新颖的线性注意力架构,在各种基准测试中取得了最先进的性能。系统展示了持续优化能力,在自主研究过程中不断提高架构质量。性能指标显示基准分数稳步提高,损失值持续降低,综合适应度分数显示出快速的初始改进,随后逐渐趋于平稳。
6. 总结
ASI-Arch代表了人工智能在自主科学发现领域的重大进步。通过结合多智能体系统、检索增强生成和进化算法,该框架能够自主进行从假设生成到实证验证的完整科学研究周期。其模块化设计和清晰的架构使其易于扩展和维护,为未来的自主研究系统提供了一个强大的模板。
项目的三个核心组件——自主架构发现管道、架构数据库和认知基础——协同工作,形成了一个完整的闭环系统,能够不断学习和改进,为人工智能驱动的科学研究开辟了新的可能性。