How to train multiple models with a shared loss and separate optimizers in MLX?

[cshan]

Hi,

I have a training scenario where multiple models share a single loss. I’ve implemented a custom algorithm that computes the gradients for all model parameters in a single pass (for performance reasons), and I’d like to update each model using its own optimizer.

I’d like to ask:

Does MLX support this kind of training setup?

What is the recommended way to implement this in MLX?

Are there any plans to support higher-level abstractions for multi-model training with separate optimizers?

Thank you for the great work—MLX is a pleasure to use!

[awni]

Yes MLX definitely supports using multiple optimizers. The simplest way is probably with the MultiOptimizer class.

Here's a pretty bare bones example which gets at how to do this:

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim

o1 = optim.SGD(learning_rate=1)
o2 = optim.SGD(learning_rate=2)

opt = optim.MultiOptimizer([o1, o2], [lambda p, _: "model1" in p])

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.model1 = mx.array(1.0)
        self.model2 = mx.array(2.0)

model = Model()

grads = {"model1": mx.array(1.0), "model2": mx.array(2.0)}
opt.update(model, grads)

[cshan]

Thank you. I found that there is no MultiOptimizer in MLX Swift yet, maybe I need to wait for a new version.

[awni]

Yes.. an issue on the mlx-swift repo would be helpful so we know to prioritize it.