Enhancing MLOps efficiency through distributionally invariant models

MLOps has emerged as a critical challenge in the field of artificial intelligence, due to the necessity for continual model updates. This requirement arises from common occurrences such as model degradation, shifts in input data distribution and changes in incidence rates. A significant bottleneck in these automated updates is the drift in the output score distribution that requires incremental effort from model consumers to change their operating points or refresh their downstream models, a cost that must be borne with each model update. Existing works address these issues via score calibration or learning an independent mapping function between the new distribution and the old one. We propose to generate distributionally invariant models at the time of model training using additional constraint of distribution matching in the loss function. We showcase the superiority of using Wasserstein distance for distribution matching over the commonly employed KL divergence. Experimental results across datasets demonstrate that the proposed approach maintains the positive prediction counts with an average recall lift of 100 bps and is versatile under challenging situations of high distribution drifts, imbalanced datasets and diverse model architectures.