The true test dataset for a generalised task

A new approach to test datasets for generalised machine learning tasks proposes drawing the test set from a distribution as different as possible from the training and validation sets, while still being in-distribution for the task itself. This method, illustrated with examples like training on Python and testing on Java, aims to measure task proficiency rather than overfitting to a specific data distribution. The idea is supported by prior work such as SCAN, DomainBed, and WILDS benchmarks.

The standard ML pipeline is well known: you want to teach a task to an algorithm. Take your data which encodes that task, split it into training, validation, and test sets. Train the ML on the training set, using the validation set to tune hyper-parameters. Then, finally, check that the algorithm has learnt the task by evaluating it on the test set. This will ensure that it hasn’t overfitted to the training data. Particularly important: the test set must come from the same distribution as the training set, or else its test is useless. That three-way split works very well – as long as the task is “process any data drawn from a similar distribution in a similar way”. For a more general task, such as “describe how a program would run, from its code”, one should draw the test dataset from a very different part of the distribution. For example, train and validation could contain Python code, and the test might be entirely in Java. Then, if the test is a success, we would be confident that the task really is “describe how a program would run, from its code” and not “describe how a very specific subset of Python programs would run, from their code”. To reach that objective, the train and validation sets could be diversified themselves – maybe train is Python and C++, validation is Python, C++, and Lisp. But in any case, the test set should be as out of distribution as possible from the train and validation set, while still in-distribution for the task. That way a successful outcome is much stronger evidence of task proficiency. And what is “in-distribution to the task”? Well, that’s whatever the programmer wants it to be. If the task is truly “describe how a program would run, from its code” then Java is certainly in-distribution. If the desired task is merely “describe how a Python program would run, from its code”, then the test set must be in Python – but it would be ideal if it was a completely different style of Python code, doing very different things, as far as possible from the training and validation sets. This new type of test set also prevents overfitting – but overfitting of a different kind. Traditionally, test sets prevent overfitting to the data sample, while trying to fit to the distribution it is drawn from. But this type of test set measures overfitting to that very distribution, rather than to the task itself. One caveat: this test set is single-use or pretty close to single-use, since repeatedly tuning against it turns it into just another validation set, overfitting to the new distribution rather than the task. Ideally there would be several distinct OOD test sets in reserve and they would be used sparingly. By no means exhaustive, but to show that ideas like these are known in the field: SCAN Lake & Baroni, 2018 — models that aced random test splits of a compositional command language failed badly on splits held out by DomainBed Gulrajani & Lopez-Paz, 2020 — the standard formalisation of leave-one-domain-out test; notably argues that model selection is non-trivial once the test distribution differs from training. WILDS Koh, Sagawa et al., 2021 — a benchmark suite of real-world tasks tumours across hospitals, wildlife across camera traps shipping with official out-of-distribution splits, so scores measure task proficiency rather than fit to the data's collection conditions.