We Raised $8M Series A to Continue Building Experiment Tracking and Model Registry That “Just Works”
"I would say the main argument for using Neptune is that you can be sure that nothing gets lost, everything is transparent, and I can always go back in history and compare."
ailslab is a small bioinformatics research group on a mission to make humanity healthier. That mission is to build models which might someday save your heart from illness. It boils down to applying machine learning to predict cardiovascular disease development based on clinical, imaging, and genetics data.
Research is so intense that it required a custom infrastructure (which took about a year to build) to extract features from different types of data:
- Electronic health records (EHR)
- Time-to-event (regression methods)
- Image (convolutional neural networks)
- Structured data and ECG.
With a fusion of these features, precise machine learning models can solve complex issues. In this case, it’s risk stratification for primary cardiovascular prevention. Essentially, it’s about predicting which patients are most likely to get cardiovascular disease.
- 1Define the task to be solved (e.g., build a risk model of cardiovascular disease).
- 2Define the task objective (e.g., define expected experiment results).
- 3Prepare the dataset.
- 4Work on the dataset in interactive mode with Jupyter notebooks; quick experimenting, figuring out the best features for both the task and the dataset, coding in R or Python.
- 5Once the project scales up, use a workflow management system like Snakemake or Prefect to transform the work into a manageable pipeline and make it reproducible. Without that, it would be costly to reproduce the workflow or compare different models.
- 6Create machine learning models using PyTorch Lightning integrated with Neptune, where some initial evaluations are applied. Log experiment data.
- 7Finally, evaluate model performance and inspect the effect of using different sets of features and hyperparameters.
“Neptune works flawlessly, and integrating it with PyTorch Lightning was very smooth.”
5 problems of scaling up Machine Learning research
ailslab started as a small group of developers. Collaboration became more challenging, and new problems began to appear along with the inflow of new team members. They noticed those problems quickly, and Neptune helped solve them.
- Data privacy
- Workflow Standardization
- Feature and Model Selection
- Experiment Management
- Information Logging
Due to limited data, ailslab has to work with NDA-protected private data. Models can only be trained locally, and the data can’t be uploaded anywhere, so there was a lot of additional work to make sure the data was safe.
ailslab separated the data workflows from the analysis workflows with Neptune in between.
Neptune keeps sensitive data safe, it just receives logged information that researchers decide to share, and the training part can happen on a local machine or anywhere else. It gives ailslab maximum control.
As a small team, ailslab didn’t need standardized practices. But as the team kept growing and new developers brought along different programming styles, it became harder to manage code development.
Thanks to Neptune, researchers use a standard library to build models, which is much easier than writing custom code. In addition, they have a standardized view of logged information with the PyTorch Lightning integration. All team members use the same infrastructure.
Neptune unifies how everyone presents results, so there’s less miscommunication.
Detailed patient records consist of multimodal data, which means multimodal training for ailslab researchers. It’s a highly complex process with an enormous amount of moving parts to test. With as many experiments as the ailslab team does, it was difficult to keep track of all models trained on different versions of a dataset and its features.
Selecting the best features for their models is easier because comparing them in Neptune is quick and straightforward. Researchers compare model performance, detailed parameters, and even hardware resource consumption. Scale is no issue. Neptune handles any amount of experiments that ailslab researchers throw at it.
ailslab researchers had to do manual tasks like creating checkpoints by hand or figuring out how to change one or more hyperparameters to do another experiment.
ailslab leaders can supervise researchers and compare their experiments all in one dashboard. The team no longer cares about organizing experiments, as Neptune does it quite elegantly. Plus, Neptune versions data for better control of experiments.
With a custom logger, it was challenging to answer essential questions when analyzing experiments. Plus, a custom logger comes with the burden of managing the logger long-term and adding new features when necessary. When bugs happened, that’s even more time to build internal tools instead of doing research.
Neptune can group experiments for comparison. It’s easy to get a link to share the results with another researcher or stakeholder.
Even if a researcher leaves the project and is no longer available, Neptune saves all information about their experiments.
Neptune automatically logs each experiment through the API. All experiments are visible across all members of the team, which makes the whole project transparent.
Why ailslab chose Neptune
In short – because it saves time. If you’re a researcher, you know that managing multiple experiments is challenging. With such complex objectives and workflows, the ailslab team has to do a lot of tedious work to stay on the right track.
Neptune saves time by removing a lot of that tedious work, and time is a luxury that the ailslab team doesn’t have a lot of!
“So I would say the main argument for using Neptune is that you can be sure that nothing gets lost, everything is transparent, and I can always go back in history and compare.”
- Compared to using a custom logger, Neptune takes care of everything, and the team has more time to do research tasks.
- ailslab researchers now use one platform with their results presented in the same way. It leaves less room for mistakes.
- Comparing and managing experiments takes less time. Researchers can go back and forth between the history of experiments, make changes, and see how the changes affect the results
- Building complex models (like deep learning models for images) and exploring how they work is a bit easier. Neptune stores data about the environment setup, the underlying code, and the model architecture.
- Neptune helps organize things. In ailslab, they add experiment URLs from Neptune to cards in their Kanban board in Notion. This easy access to experiment information helps keep everything organized. The whole team has a better idea about things like the effect of hyperparameters on the model.
All in all, machine learning is hard. Building ML models to detect heart disease before it happens adds another very thick layer of difficulty. We’re glad that Neptune takes away the tedious parts of ailslab projects, and we wish them all the best in their research.
To learn more about ailslab, check out their full story.
Want to seamlessly compare and manage experiments?
- Industry Bioinformatics, Medical
- Location Berlin, Germany
- Team size >10
- Frameworks Jupyter Notebooks, PyTorch Lightning, Snakemake, Prefect
- Neptune use cases experiment tracking, experiment management