Jane Street Software Engineer Interview Experience — Maximizing Reward as Pixels Reveal One Row at a Time

The interview was about solving an actual problem. The setup was an image classification problem using the FashionMNIST dataset. The model (a plain CNN) and the training code were already written, and the task was to modify it question by question from there; the environment they gave me was a Jupyter notebook (.ipynb).

There were mainly two parts. After understanding the code (there wasn't much to understand), the new task was that the image's pixels pop up row by row (before popping up, the pixels are black), and then you tabulate the accuracy of their already-trained model on the image as the number of visible pixel rows increases (they first gave me plotting code and asked what the resulting plot means — which is basically the above). You can see that as more pixels become visible, accuracy increases. Then they asked: based on the plot above, if you define a reward like this — for a masked image, the reward is 0 if the classification is wrong, and the reward is however many pixels remain unmasked if the classification is correct — then for this model, in order to maximize the reward, if you had to pick one universal number of pixels to unmask, how would you pick it (you basically just sweep through the data above). Question: how do you improve the model's reward (the naive idea: augmentation — mask out some pixels during training), and implement it. Then they asked: how do you improve it further (for example, when doing the augmentation, the random masking of the image can't be too random, because if you mask too much — say only one row is left visible — it might not have much effect; here I said I'd tune the range some more). Then they asked: if you could only retrain the model twice, how would you tune it? The second part was: given a trained version of the model above, suppose that for every image, the pixels pop up one at a time — the model has to dynamically choose an appropriate moment to output a result for each image. What strategy would maximize the reward? In other words, you're given, in sequence, versions of the same image where only one pixel is visible, then two pixels, then three pixels... all the way up to all pixels — and at some point you need to decide to output right then (you can forward-pass on every version without outputting anything). What strategy would you use? The strategy I proposed was: when the n-th pixel becomes visible, look back over the previous 50 steps (for example — a sliding-window idea) and count how many of them had the current highest-probability class as the highest-probability class; if the current top class was also the top class in, say, 30 out of the previous 50 steps, then output it.