Self-Healing Machine Learning: A Framework for Autonomous Adaptation in Real-World Environments

Paulius Rauba; Nabeel Seedat; Krzysztof Kacprzyk; Mihaela van der Schaar

NeurIPS 2024 · interactive exposition · paper #02

Self-Healing
Machine Learning

Paulius Rauba · Nabeel Seedat · Krzysztof Kacprzyk · Mihaela van der Schaar · University of Cambridge

Existing concept-drift methods are reason-agnostic — they pick from a fixed menu of fixes without ever asking why the model is failing. We instead frame adaptation as an optimization problem over actions, then build an agentic system $H -LLM$ that uses an LLM to diagnose the structure of the shift and prescribe a corrective action that targets the actual cause.

round

0 / 120

policy · deployed

shift type

shift severity · 0.40

space play · r reset · 1-4 shift type

baseline acc

94.2%

reason-agnostic

ℋ-LLM acc

94.2%

Δ +0.0 pp

panel 01

Live deployment under shift

A simple binary classifier was fit on a reference sample drawn from $P_{ref}$ . Each round, fresh data arrives from the shifted DGP $P_{t}$ — the dotted boundary is what the model learnt; the live points are what it sees now.

class 0 class 1 · training boundary ℋ-LLM after fix

drift score · $KL (P_{t} ∥ P_{ref})$

0.00 / thresh 0.28

expected risk under $P_{t}$

0.06 vs ref 0.06

samples observed

batch size 80 · 120 rounds total

panel 02

The ℋ-LLM loop · diagnose → prescribe → verify

When the drift score crosses threshold, the agent is invoked. It inspects the structure of the change, hypothesizes a cause, and chooses an action whose mechanism actually targets that cause — rather than the one-size-fits-all retrain of reason-agnostic methods.

ℋ-LLM · diagnostic stream

// idle · waiting for shift signal …

panel 03

Reason-aware vs reason-agnostic adaptation

The baseline is a textbook drift-adaptation policy: detect → unconditional retrain on the most recent window. ℋ-LLM does the same loop, but chooses an action conditional on the diagnosed cause. The win comes from not retraining when retraining would hurt — and from using the right kind of retrain when it would help.

baseline · reason-agnostic ℋ-LLM · reason-aware reference (no shift)

section 04

How it actually works

We formalize self-healing as an optimization over an action space $A$ . The agent picks the action that minimizes expected risk under the (now shifted) DGP — but conditional on its current diagnosis.

Let $D_{t} \sim P_{t}$ denote a batch from the live data-generating process at time $t$ . A self-healing system is a policy $π : D_{1 : t} \to A$ choosing an adaptation action so as to minimize:

a_{t}^{⋆} = ar g a \in A min E_{(x, y) \sim P_{t}} [ℓ (f_{a} (x), y)]

Reason-agnostic baselines collapse this to a fixed action (say, "retrain on the last $W$ samples"). They never ask why performance dropped, and so they can't choose actions whose mechanism matches the actual shift.

ℋ-LLM instead splits the policy into two LLM-guided stages:

diagnose · propose a hypothesis $h_{t}$ about the structure of $P_{t} - P_{ref}$
prescribe · pick $a_{t} \in A$ whose mechanism targets $h_{t}$
verify · accept the action only if held-out risk improves

The catalogue $A$ here includes refit_recent, importance_reweight, drop_spurious_feature, recalibrate_threshold, and no_op — each more or less appropriate depending on whether the diagnosis is covariate, label, concept, or spurious.

class HLLM:
    def step(self, batch_t):
        # 1 · detect
        d = drift_score(batch_t, self.ref)
        if d < self.tau:
            return "no_op"

        # 2 · diagnose (LLM reasons over shift signature)
        sig = shift_signature(batch_t, self.ref)
        h_t = self.llm.diagnose(sig, ctx=self.task_ctx)

        # 3 · prescribe (action picked to target h_t)
        a_t = self.llm.prescribe(h_t, actions=ACTION_SET)

        # 4 · verify on held-out fold
        f_new = apply_action(self.f, a_t, batch_t)
        if risk(f_new, holdout) < risk(self.f, holdout):
            self.f = f_new
            return a_t
        return "reject(" + a_t + ")"

section 05

Cite

@inproceedings{rauba2024selfhealing,
  title     = {Self-Healing Machine Learning: A Framework for Autonomous Adaptation in Real-World Environments},
  author    = {Paulius Rauba and Nabeel Seedat and Krzysztof Kacprzyk and Mihaela van der Schaar},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2024},
  url       = {https://openreview.net/forum?id=f63DKIpx0I}
}