Choose batches manually

[ibengtsson]

Hi!

Is it possible to control how the batches are chosen? I'm working on a custom loss function where it would be advantageous to control how the dataset is being batched.

[MilesCranmer]

Interesting! Unfortunately there’s nothing built in. You could add this yourself in a forked version though? Not sure if you are using Python or Julia here. On the PySR docs there’s a page explaining how to customise the backend version.

The batching is just done in this line: https://github.com/MilesCranmer/SymbolicRegression.jl/blob/cd45719835fdb5260defa7b9bb5cfa0b315eaa4a/src/Dataset.jl#L304. So you could modify that rand call to do something else?

[ibengtsson]

Thanks for pointing me in the right direction, I'll give it a go!

[ibengtsson]

I'm not really fluent in Julia but I think I managed to get a dirty fix working by modifying the line you pointed me towards like this:

function batch(rng::AbstractRNG, dataset::BasicDataset{T,L}, batch_size::Int) where {T,L}
    # If batch assignments are provided, use them
    if dataset.batch_assignments !== nothing
        # Randomly select a batch
        unique_batches = unique(dataset.batch_assignments)
        batch_id = rand(rng, unique_batches)
        
        indices = findall(x -> x == batch_id, dataset.batch_assignments)
        return SubDataset{T,L,typeof(dataset),typeof(indices)}(dataset, indices)
    else
        # Fall back to original random sampling behavior
        return batch(dataset, rand(rng, 1:(dataset.n), batch_size))
    end
end

Then I just passed a batch_assignments vector (basically a vector of length n_datapoints with integers labelling batches) via PySRRegressor.fit() and made sure it was added to all necessary functions in the Julia backend.

From an API-perspective I don't think this is the best approach though, as the batch_size argument is overridden when a batch_assignments vector is used.

[MilesCranmer]

Looks ok to me! As long as it works for you. Though one concern is it needs to compute unique(dataset.batch_assignments) every single batch. I guess you could precompute those in the BasicDataset constructor and store them in a dataset.unique_batch_assignments.

For a general approach in the future implemented in the public API I might take a slightly different approach but I will wait until it is requested a few more times – to gauge interest – before moving forward with such a feature in PySR itself.

[gm89uk]

@MilesCranmer, relevant to this, I have been experimenting on batching on and off Although with batching on, it convergences much quicker, it seems to hit an evolutionary dead-end earlier and plateau. In comparison, disabling batching (although the runtime is much slower), achieves a lower loss function. Essentially, no batching has a higher ceiling for a lower loss function if both ran indefinitely.

Although this is partially mitigated by increasing the batch_size, but a random batch will under sample the extremes of data.

I think adding more sophisticated batch samplers would bridge this gap.
Random would be O(1) complexity per batch, where as a Sobol or Latin Hypercube Sampling (LHS) might perform better? However, it could be expensive with O( * <num_x_vars>) complexity per iteration.

You could have something like batch_method = random, as default, or sobol or even a custom_batch_func as you do with other parts in the API.

Alternatively you could consider having a batch_sequence = for a more expensive batch method like a farthest point sampling which would just cycle through the batch_sequence for each iteration and leave it to to the user to generate the batch_sequence before the run.