Table of Contents
Fetching ...

ML-On-Rails: Safeguarding Machine Learning Models in Software Systems A Case Study

Hala Abdelkader, Mohamed Abdelrazek, Scott Barnett, Jean-Guy Schneider, Priya Rani, Rajesh Vasa

TL;DR

This paper targets the robustness gaps in production ML within software systems, focusing on safety, security, transparency, and silent failures. It introduces ML-On-Rails, a protocol that combines a comprehensive ML Safeguards suite with a Model-2-Software communication protocol to enforce model specifications and enable explicit error signaling. The approach includes adversarial attack detection, out-of-distribution detection, explainability via SHAP, input validation, and NeMo Guardrails for generative AI safety, all orchestrated through a runtime guard sequence. The authors validate the framework with a MoveReminder case study, demonstrating how safeguards interact across an activity-recognition model and a GPT-3.5 LLM to produce personalized recommendations while signaling and preventing unsafe or out-of-distribution results. Overall, ML-On-Rails aims to provide a practical, extensible blueprint for building robust, trustworthy ML-enabled software systems in production.

Abstract

Machine learning (ML), especially with the emergence of large language models (LLMs), has significantly transformed various industries. However, the transition from ML model prototyping to production use within software systems presents several challenges. These challenges primarily revolve around ensuring safety, security, and transparency, subsequently influencing the overall robustness and trustworthiness of ML models. In this paper, we introduce ML-On-Rails, a protocol designed to safeguard ML models, establish a well-defined endpoint interface for different ML tasks, and clear communication between ML providers and ML consumers (software engineers). ML-On-Rails enhances the robustness of ML models via incorporating detection capabilities to identify unique challenges specific to production ML. We evaluated the ML-On-Rails protocol through a real-world case study of the MoveReminder application. Through this evaluation, we emphasize the importance of safeguarding ML models in production.

ML-On-Rails: Safeguarding Machine Learning Models in Software Systems A Case Study

TL;DR

This paper targets the robustness gaps in production ML within software systems, focusing on safety, security, transparency, and silent failures. It introduces ML-On-Rails, a protocol that combines a comprehensive ML Safeguards suite with a Model-2-Software communication protocol to enforce model specifications and enable explicit error signaling. The approach includes adversarial attack detection, out-of-distribution detection, explainability via SHAP, input validation, and NeMo Guardrails for generative AI safety, all orchestrated through a runtime guard sequence. The authors validate the framework with a MoveReminder case study, demonstrating how safeguards interact across an activity-recognition model and a GPT-3.5 LLM to produce personalized recommendations while signaling and preventing unsafe or out-of-distribution results. Overall, ML-On-Rails aims to provide a practical, extensible blueprint for building robust, trustworthy ML-enabled software systems in production.

Abstract

Machine learning (ML), especially with the emergence of large language models (LLMs), has significantly transformed various industries. However, the transition from ML model prototyping to production use within software systems presents several challenges. These challenges primarily revolve around ensuring safety, security, and transparency, subsequently influencing the overall robustness and trustworthiness of ML models. In this paper, we introduce ML-On-Rails, a protocol designed to safeguard ML models, establish a well-defined endpoint interface for different ML tasks, and clear communication between ML providers and ML consumers (software engineers). ML-On-Rails enhances the robustness of ML models via incorporating detection capabilities to identify unique challenges specific to production ML. We evaluated the ML-On-Rails protocol through a real-world case study of the MoveReminder application. Through this evaluation, we emphasize the importance of safeguarding ML models in production.
Paper Structure (8 sections, 3 figures, 2 tables)

This paper contains 8 sections, 3 figures, 2 tables.

Figures (3)

  • Figure 1: The proposed ML-On-Rails protocol. We propose that ensuring robustness of ML-enabled systems requires safeguards for input validation, security; exemplified by the adversarial defence rail, safety; exemplified by OOD detection rail, and explainability. It is important to highlight that the proposed protocol components is a suggested design, providing developers with the flexibility to enable, configure, disable specific guards, or introduce additional guards. The outlined sequence of guard execution starts upon receiving an HTTP request from a client application. The process involves validation of the input against model requirements. Subsequent checks are conducted for adversarial defence and OOD detection. Once these checks pass, the model inference is performed, followed by the execution of the explainability guard.
  • Figure 2: The flow of the MoveReminder application with the proposed ML-On-Rails protocol incorporated.
  • Figure 3: Example output of the MoveReminder application with ML-On-Rails. The equipped OOD detection safeguard identified the issue and triggered a server-side error labeled as "OUT_OF_DISTRIBUTION."