Enhancing mAIstro
Making Small Language Models Reliable in Multi-Agent Medical AI Workflows
FAST-NUCES Karachi · supervised by Ms. Alina Arshad · in collaboration with the original mAIstro authors Dr. Eleftherios Tzanis and Dr. Michail E. Klontzas
Abstract
mAIstro is an open-source multi-agent framework that runs exploratory data analysis, radiomic feature extraction, medical image segmentation and classification from natural-language requests. It succeeds on every task with frontier models such as GPT-4o and Claude Sonnet 3.7, but cloud-hosted models conflict with data-sovereignty rules (HIPAA/GDPR), cost a lot and add latency. When moved to locally deployable small language models such as Llama 3.1 8B or Mistral 7B, task success drops to 10 to 18%, leaving resource-constrained clinics without agentic AI.
We close this gap by targeting the root causes of small-model failure in agentic pipelines: a decoupled architecture of 13+ specialized agents; JSON-driven tool execution with schema-constrained configuration in place of fragile code generation; a four-stage optimization pipeline of context pruning, enforced chain-of-thought, constrained decoding with dynamic schema validation and parameter sanitization; agentic RAG with dynamic schema slimming over a 10,000+ line tool library; and CPU-optimized tool defaults. The result is a fully offline, privacy-preserving pipeline that reaches 91% task success across four open models on modest hardware.
The problem
Built for frontier models, broken on small ones
The original system routes every request through one Master Agent and eight sub-agents that generate Python code. Small models overflow the context, lose the thread of reasoning, invent tools and write code that does not run.
Method
Engineering reliability back in
13+ agents with training and inference split apart, so no agent faces an oversized tool set that invites routing hallucinations.
Schema-constrained tool calls replace generated Python, removing syntax errors and invented libraries.
Context pruning, enforced reasoning, constrained decoding and parameter sanitization, each fixing one failure mode.
FAISS retrieval with schema slimming returns only the tool documentation a step needs.
CPU-optimized defaults for PyCaret, TotalSegmentator and nnU-Net run without a GPU.
No patient data leaves the machine, and no cloud infrastructure is needed.
Results
Four open models, one level of reliability
Each model was evaluated on 11 task categories, against the original paper's baseline.
| Model | Params | Baseline | Optimized | Change |
|---|---|---|---|---|
| Mistral 7B | 7B | 10% | 91% | +81 |
| Llama 3.1 8B | 8B | 18% | 91% | +73 |
| Mistral 3.1 24B | 24B | 36% | 91% | +55 |
| Llama 4 Scout | 17B | 91% | 91% | 0 |
Task success rate in percent; change in percentage points. Every model passes 10 of 11 task categories after optimization.
Runs on consumer hardware
| Component | Unoptimized | Optimized |
|---|---|---|
| TotalSegmentator | ~45 min / scan | ~10 min / scan |
| PyCaret training | Unbounded | Under 10 min, budget-capped |
| PyTorch training | Variable | Early stopping at convergence |
| Prompt size | Baseline | 15–30% smaller |
Measured on a standard CPU with no GPU.
System
No programming required
The original mAIstro authors, Dr. Michail E. Klontzas and Dr. Eleftherios Tzanis, held planning meetings with our team and are independently evaluating this system ahead of a joint MICCAI 2027 submission.
More research