MLGuard: Defend Your Machine Learning Model!

Robustness in business software where the domain is well understood is achieved through software testing, and adherence to best practices and processes. However, for Machine Learning (ML) systems this is insufficient. ML systems are dependent on data input streams that are non-stationary. As a result, ML is behaviour is underspecified (D’Amour et al., 2022) in the presence of subtle changes in the data (i.e. data shift (Moreno-Torres et al., 2012)). Data schema validation alone is insufficient as detecting violations of these conditions, e.g., out of distribution data (Hendrycks and Gimpel, 2017), can only be done probabilistically. We hypothesise that robustness can be incrementally realised in the context of ML through an interface specification (contract) that a) operates on point-estimates and distributions, b) encapsulates modelling assumptions, and c) models uncertainty as a first class citizen.

To achieve robustness for ML, research has focused on testing against noisy or malicious input data (Applis et al., 2021; Gao et al., 2020; Zhang, 2020; Kwiatkowska, 2020). However, the world is non-stationary and assumptions made offline are not guaranteed to hold during system operation. In addition, runtime validation is required to handle ML specific failure modes (i.e. miscalibration (Guo et al., 2017) and performance degradation under data shift (Moreno-Torres et al., 2012)) independent to how the model responds to adversarial examples. Our research provides the constructs and plumbing for software engineers to leverage existing algorithms that detect these failure modes (due to the difficulty of detecting these failure modes, algorithms in turn may require additional ML models to perform the validation).

Studies have managed to produce data validation tools for ML systems that continuously monitor data that are fed into the system. Data Linter (Hynes et al., 2017) acts as simplified and general validation tool that can be used to automatically inspect data and detect miscoding lints, outliers and scaling issues (Hynes et al., 2017). Data validation tools developed in-house like Deequ (Schelter et al., 2018, 2019) and Google’s TFX allow user-specified constraints (Schelter et al., 2018; Baylor et al., 2017). Most of the validation tools available use constraints and conditions that are: deterministic in nature, provide partial support for ML failure modes, and typically act as a warning system for data with anomalous or invalid characteristics. In addition, data validation tools presented are insufficient in producing a reliable and robust ML system, since the “correctness” of data is variable depending on the context and scenario of the ML system.

Best practice recommends setting up alerting and monitoring infrastructure (Breck et al., 2017; Schelter et al., 2018; Baylor et al., 2017) and ML specific architectural tactics (Cummaudo et al., 2020). However, how to specify the conditions to monitor and the actions to take when a condition fails is left to the developer. ML toolsuites for validation (e.g., Tensorflow Data Validation) provide a set of validators and abstractions for their specification. However, these tools only offer partial support for ML failure modes and automation. Specifically, Tensorflow supports data schema and constraint suggestions, but leave more sophisticated validation checks such as out of distribution detection to the developer to implement and configure.

Inspired by concepts from design by contract (Meyer, 1992), we propose a new approach, MLGuard, for specifying and validating ML contracts. Our approach takes an ML Contract Specification and generates an ML wrapper with both the code and trained ML models for validating ML contracts. We propose an ML contract specification language with i) ML specific concepts (e.g., ‘uncertain’) and ii) actions to take when the ML contract fails (e.g., log warning, throw exception, propagate uncertainty). To the best of our knowledge, this is the first proposed approach for specifying and validating probabilistic contracts for ML. Although no formal guarantees can be made of the absolute safety of the ML system, our approach provides a structured semi-automated way to help developers work towards improving the safety of the ML applications they develop by automatically detecting and responding to contract violations. MLGuard is designed to provide the scaffolding for specifying and validating contracts that, optionally, include additional validation algorithms i.e. a data drift detector.

To motivate our work, consider the case of software for automated epileptic seizure detection—classifying a segment of electroencephalogram (EEG) data as seizure or non-seizure. While a large number of papers have developed ML models for this task (Shoeibi et al., 2021), these models come with (often undocumented) conditions that must be satisfied for the output produced by the model to be reliable. This poses a safety concern as over-interpretation of EEG data (false positives) might lead to incorrect diagnosis and treatment, and this causes numerous medication side-effects, driving restrictions, increased chance of mental illness and discrimination of job opportunities while under-diagnosis (false negatives) causes delayed treatment which might increase the risk of mortality and other physical injuries.

We elaborate on these challenges below, and propose a vision in section 3 to address them, thereby improving the safety of ML models when deployed in real-world applications.

No machine-checkable ML specification language. Well-designed software components document inputs, outputs, types, pre/post conditions to be satisfied, and exceptions that may be raised. However, ML models lack a full machine-checkable specification. For example, even though an ML model accepts a vector with the same type and dimensions as the EEG data, this does not necessarily mean that the model is a suitable choice. To determine if the model is compatible with the application, one also needs to consider whether the statistical characteristics of the training data and modelling assumptions match those of the application domain.

To assist in assessing the suitability of an ML model, proposals have been made for standardised documentation templates (i.e. data sheets (Gebru et al., 2021) and model cards (Mitchell et al., 2019)). However, documentation templates require users of the ML model to manually read and interpret the ML model documentation (if any) without providing any machine-verifiable rules for safe use.

No mechanism to express uncertainty in validation rules. Different electrode placements, sampling frequencies, and filtering are possible. If these do not match those of the data the ML model was trained on, the ML model will still produce a result, but it cannot be trusted. The patient demographics may also affect the accuracy of the model. For example, it should not be assumed that a model trained on EEG from adults will work as well on EEG from children.

To ensure that the EEG data is compatible with the data on which the model was trained, one can make use of an out of distribution detection model to determine whether the assumptions of the model have been violated, i.e. the serving data in production should be distributed similarly to the data the model was trained on. However, this violation can only be detected probabilistically rather than with absolute certainty.

The conditions for these contracts operate in high-dimensional latent spaces. For ML, the input and learned latent spaces are where pre/post conditions are required to verify system behaviour. Currently we lack the mathematical constructs for guaranteeing behaviour across a high-dimensional latent space.

It is unclear how the application should respond to probabilistic violations. Unlike traditional software, ML behaviour is dependent on training data—demonstrating correctness offline is no guarantee of the system’s operating behaviour online. Thus violations of contracts need to be detected and responded to during operation rather than at design time. Best practice recommends setting up existing alerting and monitoring infrastructure (Breck et al., 2017). However, what is not specified is 1) how to configure alarms and alerts, and 2) how the system should respond.

We propose MLGuard for automatically validating whether incoming data conforms to an ML contract and handling violations. MLGuard is a practical approach for dealing with the limitations outlined in the Motivation. An overview of our proposed approach is presented in Fig. 1, and more detailed descriptions of each MLGuard component are provided below.

The work proposed and discussed in the previous sections led us to pose the following research questions. Our plan to answering them is further discussed in section 5.

• •

  RQ1 What are the abstractions required for specifying ML contracts?

• •

  RQ2 What software architecture is required to enable the generation of a wrapper for enforcing ML contracts?

• •

  RQ3 How effective is an ML contract in practice?

Our research will progress in three phases. Phase 1 will focus on extracting concepts for the ML contract meta-model from the literature and defining the ML contract specification language. Phase 2 will involve experimental evaluation of the the ML Contract. Finally, Phase 3 will investigate the effectiveness of MLGuard in an industry context.

Phase 1: ML Contract meta-model and specification language: To answer research question RQ1: What are the abstractions required for specifying ML contracts? we will expand an existing specification language. We plan to follow an iterative approach inspired by Grounded theory (Stol et al., 2016) to mine concepts from the literature (both academic and grey literature). The goal of this phase is to identify the concepts for specifying ML Contracts and defining a validation plan. The expected outcome from this phase will be 1) the ML Contract meta-model, 2) an ML Contract specification language, and 3) a set of ML specific conditions for validation.

Phase 2: Experimental evaluation of ML Contracts: The next phase of research will involve developing a prototype of our solution to answer RQ2: What software architecture is required to enable the generation of a wrapper for enforcing ML contracts? Concepts borrowed from Model Driven Engineering (MDE) will be applied to design the generators (code and trained models). Our experiment will evaluate ML Contracts against other automated and manual approaches to specifying validation logic for ML. The expected outcomes from this phase will be 1) a prototype tool MLGuard, 2) code templates for ML Contract wrappers, 3) configurations for training Validation ML models, and 4) a set of ML Wrappers generated for existing models.

Phase 3: Industry case study: To address RQ3: How effective is an ML contract in practice? the final phase of the study will evaluate how our approach can be integrated into existing software engineering workflows. We plan to run a series of industry case studies where practitioners evaluate MLGuard on ML projects. The focus of the case studies will be to identify i) user-acceptance of the approach, ii) barriers to adoption, and iii) ongoing maintenance implications.