Kaggle Winners Walkthroughs: Jane Street Real-Time Market Data Forecasting with Team Patrick Yam transcript

My name's Patrick and I'm a Kaggle Grandmaster since July 2023. And my main interest is more like time series data and the tabular data. And for the financial competition I um So I also got one gold medal before, which is the Juristat one, the Juristat crypto forecasting competition. So let's go into the the competition the solution. So this is just a short summary of my solution. So first for the model I use a modified axial transformer model.

That's So the only difference between my model and the standard standard axial attention model is that I replace one of the self-attention layer with the GRU layer for the time series aspect. And I For the features engineering side I didn't add too much features engineering. I only add one extra features, which is the time of day. And to train the model instead of the usual like cross-validation and doing the early stopping to pick the best model, I Here I just used the um seed averaging.

So I train the model with many different random seeds. And then I just pick pick one one epoch. And the next part is the online learning. So for the online online learning I So I basically I retrain the model every day. And then for the parameters for the optimizer I use some hyperparameter tuning package like Optuna to pick those parameters. And lastly is the inference module. So here because in Kaggle we are using the time series API and there's like a 9-hour inference time limit.

So here I implement a fast fast inference module that allows me to do the model inference for different seeds in parallel. So first is the feature side. So as I mentioned before I didn't do many features engineering here. So the only features I add is the time time features. So here I just do like simple rescale the the time time in a like a uniformly from zero to one and then I make it like Gaussian rank. So I So So here because it's uniform I don't have the rank here, but basically I transform it to a Gaussian distribution.

And for the model this is the main architecture of my model. So um First we So this is our numerical input. So first I pass pass this numerical input and do some clipping. So I just to remove some extreme values. And we also have some categorical values. So I just use the different embedding layer to transform this categorical value into a numerical um embeddings and then I concatenate them together with the numerical inputs.

I used the concatenation and the linear layer to to combine combine all these inputs together. Then I pass it to to this transformer layer 2D. I will explain it later on. This is the main module of of my model. Then I just simply used the fully connected layer to transform the hidden units to to to the model prediction. And here you can see that So the the input is actually one day of data. So I'm taking the two-dimensional data as as one one input.

So So I think like in some other For some other participants they they are using one one simple data as one one data point. Like they only care about the time series data, while in my case I'm taking the whole data as as just a one data point because my model is able to process both the time series information and the cross cross-sectional information. And in case you didn't heard about the axial attention, this is the the original axial attention layer.

So originally this is applied to image model. So the idea is very simple. Like instead of just one self-attention, you apply use two different self-attentions. So one is applied to the height axis and the other one is applied to the width axis. In our case we don't have height and width because our data is not image, but but we also have two axes. So one is the time series axis and the other one is is for the for for the axes or or samples in in our data.

So this is the the main idea for the for the two-dimensional transformer. So first you apply some some mixer to to mix the information between axes. So just like for transformer you have the token mixer, you try to use this token mixer to to to aggregate the information between between different tokens. And here we use this axis mixer to aggregate the information between different axes. And we also have time series mixer.

So So this is just to to aggregate the information for Yeah, for the the temporal information. Then we like the usual transformer model we used the fully connected layer after these two mixers. And for the mixer like you don't you don't really have to use like um self-attention for all of them. For instance like you just have to pick one of them. For example here for the axis mixer you you can either use self-attention, mean or or median median pooler.

But just to make sure that this mixer has no ordering information. So this operation should be unidirectional. Like if you change the order of different samples, the the output should should still be the same. And for the time series mixer we we can either use a self-attention with positional encoding or we can just use some recurrent neural network module like LSTM or GRU. In my case I I found that GRU works the best.

That's why I'm using my final model. And you can also use some like state model like RNN if you want. And so Yeah, so as I mentioned before the model So this model architecture is allow the allow the model to learn the the interaction between different samples as well as learn the the temporal temporal temporal information. And this is the like the a snapshot of the PyTorch model. That's That's describe how this model works.

And here I have a parameter called hour and multiplier. So the idea is The idea is just I I found that it would be useful to allow to to to allocate more resources to the time series some time series more module. So I found that if you set the model dimension to be too too large for the um for for the cross-sectional module, the model will the model will overfit quite a lot. But But actually it would be helpful if you set a bigger dimension for for the GRU GRU mode.

GRU mode. So So that's why I have this modifier that's only like Yeah, design the the the hidden unit size for for the for the time series mode. And to train the model I'm training the model with like multitask learning. So basically instead of just taking the main target and training on the I use all targets. And and I train it with different target weights. So So the way I designed the target weights is that I check the correlation.

So I try to assign So I try to assign a high weight for for those um for those target that is highly correlated with with the main target. And I train on it and get my results. And then I So after I get the results I try to ask ChatGPT like So So I just give him the the correlation matrix between different targets and I and I tell him, "Hey, I tried these weights and get these results. And can you suggest like what what what kind of target weights should I try try next?" And then he gives me some suggestions.

And then after some iteration I I designed the the final target weights. Well, actually it doesn't help me too too much. Like if I iterate more times. So like I just do this iteration for like two to three three times. And I fix the the target weights. weights at the end. And for the loss function I'm Yeah, here it's very simple. I just try to optimize the the evaluation metric for the competition directly. So it's the R squared with zero mean.

But then because here instead of a single target I have multiple targets. So um So So here what I did is I compute the R squared for each target and then I do the weight weight average among them. But for the weight average because I found that the scale of the loss for for each target are very different. Like maybe some targets is easier to predict and some some some of them is is harder. So the scales are a bit different.

So here I just do a very simple approximation that I scale each each target with the with the scale scale of the loss. So here like is a piece of PyTorch so I just do like loss divided by attention. So after that that's the gradient is gone, so this becomes just a constant. So like I'm scaling the the difference loss with with if it's a value. So a better approach is of course to scale them by gradient because like at the end when you do the optimization you care about the gradient not the not the value itself, but I didn't do it because of it's more convenience to to do it with a loss value and also personally I don't think it would change the result that much.

For the normalization I just did it like in a naive way. I just compute the global mean and deviation on the training data and then I just do the simple standardization. So here I didn't use like a rolling mean or rolling standard deviation. I just use the global one. And for for the training samples I use all data, but then I use different sample weights. Like I used this formula, so this 200 is just one one hyperparameter because like I do I don't want to assign a very small sample weights for for the data at the beginning.

But then I I I will multiply 1.5 to the sample weights for for the data with full full full sequence because I found that in the data like at some points the sequence length increase from like 800 something to 960. So give more more samples weight for those data with with a full full full sequence. And and I so when I train the model I try to make everything in in like a hyperparameter so that's you can make me easier to test the model with different parameters so I make this config file here.

And most of the parameters I I just tune them by myself, but some of them the the betas in in the optimizer I I just tune them by Optuna. So usually I what I did is I use Optuna to to um to to do some trials and then you have just picking the best parameter I will look at the slice slice plot. So like I will have a sense of how does each each hyperparameter affects the the final score. And then from the plot I will I will I will pick some I pick some hyperparameter or based on that.

And for the evaluation method I try to simulate the the private leaderboard setting. So like we have the training data and after that we have 120 days as a gap. So this is just like our public leaderboard. And afterwards we have our private private private data which is around 200 days. So the idea is I can use the train data to to train some some models offline. And then for the public leader for the public leaderboard data I can only use use online learning to to to upload my my models.

And for the private data I I will do both online learning and I will use this private data as the evaluation. And so for the simplification I found that it is it is better to to train the model with whole sample instead instead of doing cross validation mainly because I can't find a good valid cross validation scheme that is better than just hold out the last n days as the validation data. So like I also I try to like do the group k-fold like I split the data into chunks like 100 days as a chunk and then I do group k-fold, but but I found that it's like it's it's not very consistent like it doesn't align with with the public leaderboard.

That's why I didn't use it. But so there's one main disadvantage of using a whole out split is that we cannot use use the latest data to to to train our our to train our model. Because we we are using those validation data to to to pick the the best model epoch. model epoch. So instead of doing that I just so I decided to train the model with um many different random seeds and then just pick pick one one single epoch epoch for for them.

And I found that the ensemble with diff different random seeds is very effective like you can see from this plot like if you increase the the number of models the the score basically it just increases. Yeah. And and after I do all the local experiments I will retrain the model with all data compute all the like mean and standard deviation for the normalization with the full data and Yeah. And next is online learning.

So for online learning I'm doing the following like for each day I will use use Adam um use Adam optimizer to to to update the model. So basically I will do three I update the model with three three steps. And yeah you can see like this this is the value of beta one and beta two issues slightly smaller than than usual. And I also try to to just load the the the optimizer that's that's that's used in the offline training.

But for some reason it doesn't work like you you just much worse than than using a fresh fresh optimizer. Also I try to like instead of just using the current state data to train for like three three steps I I also try to get the data from like t minus one and and t minus two and train it together with with with the data in the current step, but it doesn't improve. Actually it is worse than than just using the the current state data.

And the online learning is super important here and so this is a plot to compare like with and without online learning. So the blue line is with online learning and the orange one is without online learning. So you can and the score here is a rolling 20 days average score. So you can see the model decays quite a lot after like 300 days. So here the gap is like 1% which is quite quite large. So so here like yeah basically online learning is very important in this competition.

Um and this is so the next one is the fast model inference module. So here so so because I'm training the model with many random random seeds as I showed showed you before because like basically you you got a better score if you increase the the number of model. But then the the problem becomes the inference time because we only are allowed to to like do the model inference in nine nine hours. So so here I I I what I did is is to replace all the linear model with with with this Eisen model and then I I just stack stack all the model weights.

Like for example if you have if we have 20 models in total we just have to stack all these 20 20 weights into just one thing. And after we do the stacking we can use this use this use this Eisen Eisen model that um that um can inference can inference the data at once instead of using a for loop. Also because I'm using a GRU model for the time series model so so I just have to cache the the latest hidden states instead of going from from the beginning each time.

And as the results we so it only takes me like 2.5 hours inference time for for for 17 models plus online learning. So the main problem the main bottleneck here is the online learning because for online learning you cannot do it in parallel. You still have to use a for loop for for the online learning, but for the model inference you can basically basically no matter like how many models do you have like I try to do the inference with 40 models and the inference time is the same as as you doing the the inference for for for just one model.

And so lastly this is like some findings I I have. So first is I found that the leaderboard is very aligned with with my local validation. So you can see like this is the my local validation score and leaderboard score. So like whenever I find some improvement in my local validation score it improved in the leaderboard score as well. And also the public the private leaderboard is very stable like you can see that the ranking basically it doesn't like even from from the first week.

So it's it's very stable and and yeah I must say the data quality in this competition. And yeah I think that's that's it.