[Paper Review] Aloe: A Family of Fine-tuned Open Healthcare LLMs

Ashwin Kumar Gururajan, Enrique Lopez-Cuena+11 more|arXiv (Cornell University)|May 3, 2024

Phytochemistry and biological activity of medicinal plants6 citations

TL;DR

The paper presents Aloe, a family of open healthcare LLMs fine-tuned from open base models, incorporating synthetic data, model merging, and alignment to achieve state-of-the-art performance for 7B-scale medical LLMs while addressing ethics and safety through red teaming and DPO alignment.

ABSTRACT

As the capabilities of Large Language Models (LLMs) in healthcare and medicine continue to advance, there is a growing need for competitive open-source models that can safeguard public interest. With the increasing availability of highly competitive open base models, the impact of continued pre-training is increasingly uncertain. In this work, we explore the role of instruct tuning, model merging, alignment, red teaming and advanced inference schemes, as means to improve current open models. To that end, we introduce the Aloe family, a set of open medical LLMs highly competitive within its scale range. Aloe models are trained on the current best base models (Mistral, LLaMA 3), using a new custom dataset which combines public data sources improved with synthetic Chain of Thought (CoT). Aloe models undergo an alignment phase, becoming one of the first few policy-aligned open healthcare LLM using Direct Preference Optimization, setting a new standard for ethical performance in healthcare LLMs. Model evaluation expands to include various bias and toxicity datasets, a dedicated red teaming effort, and a much-needed risk assessment for healthcare LLMs. Finally, to explore the limits of current LLMs in inference, we study several advanced prompt engineering strategies to boost performance across benchmarks, yielding state-of-the-art results for open healthcare 7B LLMs, unprecedented at this scale.

Motivation & Objective

• Promote open science in healthcare AI by developing competitive open LLMs that safeguard public interest.
• Investigate the effectiveness of instruct tuning, model merging, and alignment for open healthcare LLMs.
• Assess ethical, bias, and safety aspects through red teaming, DPO alignment, and risk assessment.
• Share training data, merging configurations, and prompting strategies to enable reproducibility and responsible deployment.

Proposed method

• Fine-tune open base models (Mistral-7B and Llama 3-8B) with domain-specific medical instruction data and synthetic CoT-enhanced data to create Aloe variants.
• Apply model merging (DARE-TIES) to combine strengths of multiple base/submodels and obtain a merged 7B-8B-size model.
• Perform two-stage Direct Preference Optimization (DPO) alignment including a red-teaming inspired dataset to improve safety and ethical behavior.
• Explore advanced prompting strategies (Medprompting, self-consistent CoT, few-shot with embeddings) during inference to boost medical benchmark performance.
• Release the DPO-aligned model under CC-BY-NC 4.0 and share merging configs and prompting repo to support reproducibility and safety analysis.

Experimental results

Research questions

• RQ1 How do instruct tuning, synthetic data, and model merging impact open healthcare LLM performance at ~7B scale?
• RQ2 Can alignment (DPO) and red-teaming improve safety, bias, and factuality in open healthcare LLMs without sacrificing medical accuracy?
• RQ3 Do advanced prompting strategies significantly boost open healthcare LLM benchmarks at this size?
• RQ4 What is the trade-off between model size, performance, and safety in the Aloe family?

Key findings

• Aloe achieves state-of-the-art results for open healthcare models at 7B when using model merging and advanced prompting.
• DPO alignment reduces unsafe responses and improves ethics and factuality metrics, though larger models may still exhibit higher bias and toxicity in some settings.
• Medprompting and self-consistent prompting with ensembles substantially boost medical benchmark accuracy, with diminishing returns beyond a certain ensemble size.
• Aloe-Llama3-8B-Alpha outperforms several larger open models and approaches the performance of bigger private models on multiple medical benchmarks.
• The paper provides public release of the DPO-aligned model and training/merging configurations to support reproducibility and safety-focused research.