128 Responsible AI: Standardising AI model training for safety and transparency using MLFlow

Background Training an AI model resembles teaching a student in problem-solving according to established rules. Similarly, determining objectives allows the AI training process to investigate various methods to identify the optimal solution. The primary aim is to develop solutions that not only performs well on familiar tasks but also demonstrates the ability to tackle new, unseen challenges effectively. This involves a combination of informed hypotheses, expert insights, and extensive experimentation. However, in practice, tracking the performance of consecutive experimentations using different configurations becomes difficult if there isn’t a clear audit trail of the experiments performed. Over time, a manual reviewing of the experiments can lead to duplication and reproducibility issues.Methods We follow a standardised process for our in-house AI model training using MLFlow, a tool setup within our secure digital space GRID. MLFlow connects the AI experiment script directly to the experiment tracking system, allowing the technical team to view all runs associated with their parameter changes. This enables direct comparison between experiments, and once satisfied with performance, models are pushed to the model registry for inference deployment.Results Since implementation, we have tracked 738 experimental runs across 10 experiments, averaging 74 runs per experiment – demonstrating thorough hyperparameter exploration. Each run automatically captures an average of 31 hyperparameters and 19 performance metrics, providing comprehensive experiment documentation. Our model registry maintains 11 production-ready models (on an average of 115 associated experimental runs monitored for further analysis) such that only the most optimized models reach the model registry. The system has also drastically improved our audit capabilities, with complete model lineage now retrievable instantly.Conclusion Integrating MLFlow into our internal AI development pipelines has standardised model training for both audit and reproducibility purposes. This has significantly improved project explainability while maintaining full transparency for safety and compliance requirements.