Since 2023, after Databricks acquired MosaicML, Mosaic AI has been incorporated into Databricks as part of its growing suite of AI and machine learning (ML) tools. But what exactly is Mosaic AI, and how can it enhance your existing ML workflows?
In this insight, we explore what Mosaic AI is, how it integrates with Databricks (including MLflow), and how to use it effectively, along with key limitations and upcoming capabilities.
1. What is Mosaic AI
In essence, the Mosaic AI framework is designed to accelerate the development, training, deployment, and monitoring of deep learning models - especially generative AI models such as Large Language Models (LLMs). It supports the classical machine learning lifecycle, but it is particularly optimised for high-performance training and fine-tuning of large-scale models.
In my opinion, the main benefit of Mosaic AI is the ability to train models more rapidly and efficiently, especially when fine-tuning open-source LLMs on your own data. These foundation models are typically pre-trained on large-scale datasets but lack domain-specific knowledge. Rather than training models from scratch, you can build upon pre-trained models using the Mosaic AI Composer, significantly reducing training time and complexity.
2. Which Mosaic AI components are integrated into Databricks
Databricks now includes several Mosaic AI-powered capabilities:
- Model Training
- Model Serving
- Vector Search
- Agent Framework
- AI Gateway
Model Training
Model training refers to the in-code integration of the Mosaic AI Composer to fine-tune LLMs or open-source models (e.g., from Hugging Face). It supports integration with MLflow for experiment tracking and evaluation. As of now, in the Western European region, this is only available via code and not yet available in the UI.
Model Serving
Models can be deployed through serving endpoints. These endpoints can be used by authorised users. Custom models must be registered in Unity Catalog, but you can also serve prebuilt foundation models available in Databricks.
Vector Search
This is specific to Generative AI use cases. It allows indexing and embedding of structured or unstructured documents, enabling you to query those documents using similarity search. Vector Search is already available in the Western European region.
Agent Framework
This supports the development of AI agents that can interact with each other or with users, allowing LLMs and classical models to communicate in workflows. However, the Agent Framework is not yet available in Western Europe.
AI Gateway
The AI Gateway enables governance and monitoring for accessing AI models. This includes features such as rate limiting, restricted access, payload logging, and compliance controls. It helps ensure safe and efficient use of models across the organisation and is already available in Public Preview, including in Western Europe.
3. The do's and don’ts with Mosaic ML
So what’s valuable and when should you use it?
- Don’t use Mosaic AI for classical non-deep learning models (e.g., regression, LightGBM, XGBoost). These models are already efficient and well-optimised for most use cases. Mosaic AI doesn’t provide significant advantages here, aside from potential GPU optimisations.
- Use Mosaic AI for:
- Fine-tuning deep learning models and LLMs, especially from Hugging Face or other open-source repositories
- Accelerating training of open-source models like BERT or LLaMA
- Scenarios that benefit from distributed training or GPU optimisation
4. Example use case
Suppose you want to fine-tune a BERT model from Hugging Face for sentiment analysis on a small tweet dataset:
- Use Mosaic AI Composer to handle the fine-tuning
- Track results and metrics via MLflow
- Register the model in Unity Catalog
- Serve the model via an endpoint using the AI Gateway
The end goal of Databricks is to provide an accessible interface for training models with Mosaic ML. This would work similarly to how AutoML is currently used, making it easy to explore the capabilities of Mosaic ML.
5. The difference between MLflow and Mosaic AI
Let’s also touch upon the difference between the Mosaic AI framework and MLflow.
MLflow focuses primarily on experiment tracking and model management. It helps you track your ML models, monitor data lineage, log experiments and runs, and manage model stages (e.g., transitioning models to staging or production). However, its functionality is mostly limited to tracking and lifecycle metadata.
Mosaic AI, on the other hand, extends beyond that - especially with the inclusion of the AI Gateway. The AI Gateway adds a governance layer on top of your models, including storing them in Unity Catalog, enabling model serving, creating custom functions, and even allowing users to ask questions about their data via those models. A big added value here is that it also provides monitoring tools to track request volumes and load, and supports fine-grained access control for securing model endpoints.
With Composer, Mosaic AI further allows you to fine-tune models on your own data, making it highly flexible for enterprise use cases.
The main added value lies in this unified and production-ready interface - combining governance, serving, tuning, and access control - which is not yet generally available, but is expected to be released in the coming months following its announcement at the Data & AI Summit 2025.
The key differences are described below.
| Feature | A/B Testing (AI Gateway) | MLflow Evaluation |
|---|---|---|
| Environment | Live production | Offline (historical data) |
| Data Type | Real-time user queries | Pre-collected datasets |
| Purpose | Compare models under real traffic | Compare models based on logged metrics |
| Key Metrics | Latency, user feedback, cost, engagement | Accuracy, precision, recall, loss |
| Decision Making | Choose the best model for deployment | Choose the best model for further tuning |
The Mosaic AI Gateway effectively addresses challenges related to quality, control, and cost, ensuring production-ready AI agent systems. A key advantage is the ability to track usage insights and intermediate results, improving observability and decision-making.
The session at the Data & AI Summit about traditional ML models at scale confirms what is said above and shows a demo of this. They use the feature store, they train models with AutoML and register them in Unity Catalog.
6. Conclusions and upcoming research
To conclude, the AI Gateway offers a scalable and governed way to serve and manage your ML models in production.
Additionally, fine-tuning with Mosaic AI Composer is specifically designed for deep learning models and large language models (LLMs), and is not applicable to traditional ML models.
Some features - such as the interface for model fine-tuning- are not yet generally available, but we plan to test them further as they become accessible (see image below). A list of features and their availability in Europe can be found here.
We're especially looking forward to exploring:
- Mosaic AI Agent Framework & Evaluation
- Foundation Model Fine-tuning UI
- Enhanced Vector Search capabilities
