I'm working on a project where I need my model to predict a sequence of n 3x3 matrices given an input sequence of n 3x3 matrices for a physical simulation. (the n's are all perfect squares e.g. I only need to consider sequences of 9, 16, 25, 36, etc. matrices all the way to 324). I need the model to approximate the mathematical relationship between the input and output (which can be computationally expensive for large n, hence the need for an ML model in the first place).

However, the problem I have is that my dataset only contains input/output values for sequences of 9, 16, ... 81 matrices (after that it becomes hard to compute).

I am currently a novice in machine learning and just know the basics of Pytorch, but so far I have trained a simple neural network with alternating Linear and ReLU layers for the n=9 case, with error ~ 10^-4. However, now I want to expand to the general case, but I am not sure of the best way to do so.

What I am thinking is to make my neural network big enough for the largest case n=324 and use some kind of RNN to make the predictions accurate for large n even though my dataset does not cover all those values.

I just wanted to ask if this is the right approach or if there is a better way to transfer the information learned from small dataset sizes to make predictions for larger cases, and also what type of RNN I should use.

Thank you!