Using Kolmogorov-Smirnov Distance for Measuring Distribution Shift in Machine Learning

TL;DR

This work addresses distribution shift between training and test data in reinforcement-learning–driven traffic signal control. It proposes using the Kolmogorov-Smirnov distance $D = \max_i |p_A(i) - p_B(i)|$ to quantify shift between traffic distributions over the 8 NEMA phases and to predict performance degradation of a DQN-based TS-C agent. Through real-world ATSPM data and SUMO-based simulations, the authors demonstrate that even small KS shifts (e.g., $D=0.02$) can substantially increase travel time and reduce throughput, with effects amplified by higher traffic volumes. The study suggests KS distance as a practical, real-time monitoring tool for deployment decisions and paves the way for comparisons with other distribution-shift metrics like Wasserstein distance in traffic domains.

Abstract

One of the major problems in Machine Learning (ML) and Artificial Intelligence (AI) is the fact that the probability distribution of the test data in the real world could deviate substantially from the probability distribution of the training data set. When this happens, the predictions of an ML system or an AI agent could involve large errors which is very troublesome and undesirable. While this is a well-known hard problem plaguing the AI and ML systems' accuracy and reliability, in certain applications such errors could be critical for safety and reliability of AI and ML systems. One approach to deal with this problem is to monitor and measure the deviation in the probability distribution of the test data in real time and to compensate for this deviation. In this paper, we propose and explore the use of Kolmogorov-Smirnov (KS) Test for measuring the distribution shift and we show how the KS distance can be used to quantify the distribution shift and its impact on an AI agent's performance. Our results suggest that KS distance could be used as a valuable statistical tool for monitoring and measuring the distribution shift. More specifically, it is shown that even a distance of KS=0.02 could lead to about 50\% increase in the travel time at a single intersection using a Reinforcement Learning agent which is quite significant. It is hoped that the use of KS Test and KS distance in AI-based smart transportation could be an important step forward for gauging the performance degradation of an AI agent in real time and this, in turn, could help the AI agent to cope with the distribution shift in a more informed manner.

Figures (8)

• Figure 1: NEMA phases for a standard 4-way intersection. Each arrow represents a traffic movement, and some phases serve multiple movements.
• Figure 2: K-S distance $D$ (black line) between a reference distribution (red) and an empirical distribution (blue).
• Figure 3: On the left, the two traffic distributions $p_A$ (blue) and $p_B$ (red). On the right, their corresponding CDFs and the K-S distance highlighted in black.
• Figure 4: Throughput ratio and Extended Travel Time for the fixed total vehicles volume and fixed phase distribution scenarios. In Figure (a), the x-axis represents the phase KS distance from the training scenario. In Figure (b), the x-axis represents the total volume of vehicles.
• Figure 5: Throughput ratio and Extended Travel Time for all combinations of total vehicle volume and phase KS distance. The x axis represents the total vehicle volume, the y axis represents the agent performance, and the different curves are colored according to their KS distance (blue = 0.0, orange = 0.1, green = 0.2, red = 0.3, purple = 0.4, brown = 0.5, pink = 0.6)