AAI_2025_Capstone_Chronicles_Combined

The advantage of training multiple models in parallel using vmap is that it sidesteps some of these issues by allowing multiple independent training runs to share the same GPU resources without requiring excessively large batch sizes for individual models. This allows researchers to balance efficient GPU usage while maintaining the benefits of smaller batch sizes for model generalization. While the ability to train multiple small models in parallel is computationally interesting, it is important to identify practical use cases where this capability is beneficial. There are several scenarios where training multiple models simultaneously is advantageous. When building an ensemble of models, training multiple networks on different bootstrap samples of the dataset can improve overall predictive performance by reducing variance. Some architectures divide computational tasks across multiple specialized submodels, requiring multiple neural networks to be trained in parallel. For example, in stock price prediction, the same neural network architecture may be applied to different financial instruments, requiring independent training runs for each stock. In reinforcement learning, multiple agents can be trained simultaneously in different environments, leading to more sample-efficient training or large numbers of agents can be simulated together in an environment to observe their interactions with each other. The development of the TurbaNet library aims to address the shortcomings of existing implementations and provide a user-friendly API for leveraging JAX's parallelization capabilities. The key objectives of this project are abstract vectorization, flexible architectures and loss functions, and efficient runtime for critical training and evaluation functions. Users should not need to manually set up vectorization for their models, the system should support any neural network structure and loss function compatible with JAX and Flax, and most importantly training should remain as efficient as possible, making full use of available GPU resources. By achieving these goals, TurbaNet will provide a flexible and efficient solution for users looking to train multiple small neural networks in parallel, filling a gap in the current landscape of machine learning tools.

130