Stochastic FW: optional batched user-provided functions

[matbesancon]

With the current SFW interface, users provide a function that processes one data point, batching happens a level higher when we call the provided functions.

One possibility would be to make users provide batched functions by default:

f_batched(θ, xs) = sum(f(θ, x_i) for x_i in xs)
g_batched(θ, xs) = sum(g(θ, x_i) for x_i in xs)

What they provide now is the equivalent of the functions f and g above.

[pokutta]

good question - i would say we leave as is for now. the reason is that we need to thing how to best map e.g., variance-reduced methods as they need special batch sizes depending on the iteration.

[matbesancon]

OK yes. Even with the alternative version, each iteration can control the batch size by picking the size of the xs list that is passed to {f/g}_batched

[pokutta]

ok i will need some extra explanation tomorrow to discuss.

[matbesancon]

So for now at the FW function level we have this:

compute_gradient(f, x, rng=rng, batch_size=batch_size)

At the compute_gradient level for f::StochasticObjective:

    rand_indices = if full_evaluation
        eachindex(f.xs)
    else
        rand(rng, eachindex(f.xs), batch_size)
    end
    return sum(f.grad(θ, f.xs[idx]) for idx in rand_indices)

So compute_gradient is the place where the default batching behaviour is defined, and calls f.grad on individual data points. The change in behaviour would be that compute_gradient passes down the rng, batch_size and full_evaluation arguments to f.grad, that is, the user-defined function, which itself implements the batching.