On-Device Model Drift Detection and Automatic Rollback for Android

Last year we hit a wall on an OCR recognition project: after a staged rollout of the on-device model, recognition accuracy dropped from 96% to 82%, user complaints poured in, and the team still spent two days just identifying the new model version as the culprit. In the postmortem, the root cause was not a bug in the model itself. The new model was extremely sensitive to the image noise distribution on low-end devices, and the training set lacked those kinds of samples. This kind of problem is called model drift.

On-device inference does not have the centralized monitoring infrastructure that server-side inference has. Once a model goes wrong, both detection and mitigation take a long chain. Below I will discuss an on-device quality assurance system I shipped, covering everything from data distribution detection to automatic rollback.

On-Device Model Drift: Why It Is Harder Than Server-Side

All server-side inference data flows through a controlled pipeline, so sampling, computing distributions, and triggering alerts can be done in one go. On-device is a completely different story.

On-device inference happens on the user’s device. Input data is never uploaded to the server, so you cannot perform real-time statistical analysis on inference inputs at a central node. Moreover, on-device environments are highly fragmented: camera modules, sensor accuracy, and compute unit instruction sets on different device models can all make the same model produce different output distributions.

I divide on-device drift into two categories:

Data Drift: the distribution of input data changes. For example, usage frequency in nighttime scenes suddenly increases, while the training set is mostly daytime samples.

Concept Drift: the input distribution stays the same, but the relationship between input and output changes. For example, after an OCR model upgrade, handwriting that the older version recognized correctly is misjudged by the new version.

The detection strategies differ between the two types, but on-device they face the same constraint: you cannot send user data back to the server for analysis. The detection logic must be pushed down to the device.

Detection Layer: Deciding Distribution Shift On Device

Choosing Statistics

Doing distribution detection on-device has a core tension: computing power. KL divergence and MMD (Maximum Mean Discrepancy) are accurate, but their computational cost is not mobile-friendly.

In practice I chose two lightweight metrics:

PSI (Population Stability Index) is used to measure the stability of model output distribution. The formula is simple:

PSI = Σ (actual_i - expected_i) × ln(actual_i / expected_i)

expected_i is the baseline distribution (the binned probabilities computed on the validation set when the model was released; the bin boundaries are also fixed in this step), and actual_i is the actual distribution within the current sliding window on the device (reusing the same bin boundaries). PSI < 0.1 means the distribution is stable, 0.1–0.25 means mild drift, and > 0.25 requires attention.

Simple feature-level statistics: instead of uploading raw data, maintain only the sliding mean and standard deviation of each input feature on the device. When uploading, report only the statistics, with no raw samples. The sensitivity is sufficient, and compliance risk is zero.

One prerequisite worth adding: output distribution PSI drift is essentially only an indirect signal. It says “the model’s recent output distribution differs from the distribution at release,” but a change in output distribution is not directly equivalent to a drop in model accuracy or the occurrence of concept drift. For example, if the user population itself changes seasonally, the output distribution will also change even though the model itself has no problem. PSI is better used as a cheap, easy-to-compute anomaly alert signal on-device. To truly determine whether model performance has degraded, you still need to combine business metrics (such as active user feedback or retry rate) or sampled manual review.

Engineering Implementation of On-Device Detection

class DriftDetector(
    private val baselineDistribution: FloatArray, // 基线分箱概率
    private val baselineMin: Float,               // 基线阶段固定的分箱下界
    private val baselineMax: Float,               // 基线阶段固定的分箱上界
    private val binCount: Int = 10,
    private val windowSize: Int = 200            // 滑动窗口大小
) {
    private val recentOutputs = ArrayDeque<Float>(windowSize)
    private var driftScore: Float = 0f

    fun observe(output: Float) {
        if (recentOutputs.size >= windowSize) {
            recentOutputs.removeFirst()
        }
        recentOutputs.addLast(output)

        if (recentOutputs.size >= windowSize) {
            driftScore = computePSI(recentOutputs.toList())
        }
    }

    private fun computePSI(samples: List<Float>): Float {
        // 关键点：分箱边界必须复用基线阶段固定下来的 min/max，
        // 不能按当前滑动窗口的 min/max 重新划分，否则基准分布和当前分布的分箱口径对不上，PSI 比较就失去了意义
        val range = baselineMax - baselineMin + 1e-7f
        val actualBins = FloatArray(binCount)
        samples.forEach { v ->
            val idx = ((v - baselineMin) / range * binCount).toInt()
                .coerceIn(0, binCount - 1)
            actualBins[idx]++
        }
        actualBins.forEachIndexed { i, _ -> actualBins[i] /= samples.size }

        var psi = 0f
        for (i in 0 until binCount) {
            val a = actualBins[i] + 1e-7f
            val e = baselineDistribution[i] + 1e-7f
            psi += (a - e) * ln(a / e)
        }
        return psi
    }
}

There is an easy trap here: if the bin boundaries are regenerated dynamically each time from the current sliding window’s min()/max(), the baseline distribution and the live distribution are effectively measured on two different scales. Even if the real distribution has not changed at all, PSI can produce a false positive because the binning criteria on the two sides are inconsistent, or conversely, mask real drift. The correct approach is to determine the bin boundaries (baselineMin/baselineMax) when the model is released and the baseline is built, freeze and persist them, and reuse this fixed set of boundaries every time PSI is computed on the live sliding window. Then the binning criteria on both sides are comparable.

Reporting Strategy

You cannot report every time drift is detected; the traffic cannot handle it. I use two threshold levels:

• Soft threshold (PSI > 0.2): record locally on the device; only report one event after it has been triggered cumulatively 3 times, with the current PSI value and device model information.
• Hard threshold (PSI > 0.3): report immediately and trigger the local degradation logic at the same time.

The reporting channel goes through the existing analytics pipeline, reusing the existing sampling rate and deduplication mechanism, so there is no need to build an additional data channel.

Rollback Layer: Model Version Management and Automatic Degradation

Detecting drift is only the first step; automatic mitigation is the core.

Version Bundle Management

On-device cannot switch traffic in seconds the way server-side can. Model files are usually 2–50MB, and downloading on demand takes time. In the app’s model management module, I maintain a “version triple”:

模型包结构:
├── model_v3.2.tflite       // 当前主模型
├── model_v2.8.tflite       // 上一个稳定版本
└── model_fallback.tflite   // 保底模型（轻量、高泛化）

The fallback model is the key to the whole scheme. It is a lower-complexity but more generalizable version. For example, replacing the MobileNet-v3 backbone with MobileNet-v1 gives faster inference and higher tolerance to input changes.

data class ModelBundle(
    val current: ModelVersion,      // v3.2
    val previousStable: ModelVersion, // v2.8
    val fallback: ModelVersion      // 保底
) {
    fun selectTarget(driftLevel: DriftLevel): ModelVersion = when (driftLevel) {
        DriftLevel.NORMAL -> current
        DriftLevel.WARNING -> previousStable  // 退回上一稳定版
        DriftLevel.CRITICAL -> fallback        // 紧急兜底
    }
}

Rollback Trigger Logic

class RollbackController(
    private val driftDetector: DriftDetector,
    private val modelBundle: ModelBundle,
    private val inferenceEngine: InferenceEngine
) {
    private var consecutiveDriftCount = 0

    fun onInferenceComplete(output: Float) {
        driftDetector.observe(output)
        val psi = driftDetector.currentDriftScore()

        when {
            psi > 0.3 -> {
                consecutiveDriftCount++
                if (consecutiveDriftCount >= 3) {
                    triggerRollback(DriftLevel.CRITICAL)
                    reportEvent("critical_drift_rollback", psi)
                }
            }
            psi > 0.2 -> {
                consecutiveDriftCount++
                if (consecutiveDriftCount >= 5) {
                    triggerRollback(DriftLevel.WARNING)
                }
            }
            else -> consecutiveDriftCount = 0 // 恢复正常，重置计数
        }
    }

    private fun triggerRollback(level: DriftLevel) {
        val target = modelBundle.selectTarget(level)
        inferenceEngine.switchModel(target)
        consecutiveDriftCount = 0
    }
}

Here is a pitfall I have stepped on: do not trigger a rollback based on a single PSI violation. Short-term behavioral fluctuations by the user, such as walking from outside into a room, can cause transient distribution changes. The consecutive count window filters out this kind of noise well.

Recovery Strategy After Rollback

After a model rollback, we cannot remain on the old version forever. After rollback I start a periodic check task: every 6 hours, run the validation script for the baseline test set with the fallback model (50 test images pre-packaged in the app). If PSI returns to the normal range, attempt to switch back to the previous stable version.

Most false rollbacks caused by environmental factors recover automatically within 24 hours. Rollbacks caused by actual model bugs wait for the server to deliver a fixed new version.

Coordination Layer: A Device-Cloud Closed Loop

On-device self-closure alone is not enough; the drift signal needs to feed back into the model iteration process.

The reporting data structure I designed:

@Serializable
data class DriftReport(
    val modelVersion: String,
    val psiValue: Float,
    val deviceModel: String,
    val androidVersion: Int,
    val chipset: String,        // 如 "Snapdragon 8 Gen 2"
    val featureStats: Map<String, StatPair>, // 特征均值/方差，无原始数据
    val driftLevel: DriftLevel,
    val timestamp: Long
)

@Serializable
data class StatPair(val mean: Float, val std: Float)

After receiving the report, the server does two things:

1. Aggregate by device model + chipset and draw a drift heat map across different device models. This can reveal model adaptation issues on certain hardware. The earlier low-end device noise-sensitivity problem was discovered through this heat map.
2. Trigger automated data backfill: feed the feature statistics corresponding to the drifting device models back to the labeling platform to supplement training samples in a targeted way.

This mechanism has been running for more than half a year. The model’s online incidents dropped from 2–3 per quarter to zero, and the false rollback rate is controlled within 5%.

Implementation Suggestions

Three suggestions from real-world practice:

Build the fallback model well before doing monitoring. Without a fallback plan, once monitoring finds a problem, there is nothing you can do about it. The fallback model should not pursue ultimate accuracy; it should pursue generalization and stability. Its inference time is best kept within 60% of the main model’s.

Drift thresholds need to be calibrated by business scenario. Scenarios with high real-time requirements, such as AR navigation, should use more sensitive thresholds, while offline batch processing, such as album classification, can be relaxed. I suggest going live first and running for a week to collect baseline data, then setting thresholds according to business tolerance.

On-device fault tolerance is more important than detection precision. Model drift detection itself can also misjudge, but the cost of rolling back to the previous stable version is far lower than continuing to use a problematic model. Focus the design on the reliability of the rollback path, and do not overinvest in the fine-tuning of the detection algorithm.