How to track the current iteration number in a custom loss function?

[GFODK]

Hi,
I'm training a symbolic regression model in Julia (not using PySR) and defining a custom loss function. My loss function includes two additional conditions, and I manually assign different weights to each term, like this:

w_a, w_b, w_c = Ref(0.1), Ref(0.4), Ref(0.4)

function custom_loss(tree::Node, dataset::Dataset{T,L}, options::Options) where {T,L}
    ...
    loss = L_a * w_a + L_b * w_b + L_c * w_c
    return loss
end

model = SRRegressor(
    niterations=1000,
    binary_operators=[+, -, *, /],
    unary_operators=[sin, cos, exp, log],
    population_size=50,
    loss_function=custom_loss
)

mach = machine(model, x, y)
fit!(mach)

I’d like to adjust the weights of my loss terms dynamically based on the current iteration.
Is there a built-in variable that keeps track of the iteration count inside the custom_loss function? Or is there another way to access the current iteration during training?

Any help would be much appreciated! Thanks! 😊

[MilesCranmer]

Hi @GFODK,

Unfortunately you cannot have a dynamic loss function, because it is cached at various points in the search rather than re-computing each time. So if your loss function changes, it will break some assumptions in the code.

What you can do though is adjust the loss function and call fit! again—it will start where it left off, and recompute all the losses on existing populations of expressions. So, for example:

model = SRRegressor(
    niterations=1,
    binary_operators=[+, -, *, /],
    unary_operators=[sin, cos, exp, log],
    population_size=50,
    loss_function=loss_functions[1]
)
mach = machine(model, x, y)
fit!(mach)

for i in 2:100
    mach.model.loss_function = loss_functions[i]
    mach.model.niterations += 1
    fit!(mach)
end

or something like this. You might want to build some callable struct that has parameters you can declare, like

struct LossContext <: Function
    w_a::Float64
    w_b::Float64
    w_c::Float64
end

function (ctx::LossContext)(tree::Node, dataset::Dataset{T,L}, options::Options) where {T,L}
    ...
    return L_a * ctx.w_a + L_b * ctx.w_b + L_c * ctx.w_c
end

and then you can just define the loss like

    SRRegressor(loss_function=LossContext(0.1, 0.4, 0.4))

[GFODK]

Hey, thanks so much for the quick response! 🙌

This is my first time working with Julia, so I hadn’t thought of that approach. I ended up rewriting my code by combining both of the methods you suggested. Now, I'm updating the loss weights (w_a, w_b, w_c) at each iteration and assigning them to loss_function.

However, I'm seeing the following warning messages during training. Do you think it's safe to ignore them, or could they indicate that something is going wrong with the training process?

Here’s the modified code—I'm updating the loss function in every iteration using a for loop:

struct LossContext <: Function
    w_a::Float64
    w_b::Float64
    w_c::Float64
end

w_a, w_b, w_c = 1.0, 1.0, 1.0
model = SRRegressor(
    niterations=1,
    binary_operators=[+, -, *, /],
    unary_operators=[sin, cos, exp, log],
    population_size=100,
    loss_function=LossContext(
        w_a=w_a, 
        w_b=w_b, 
        w_c=w_c
    )
)
mach = machine(model, x, y)
fit!(mach)

# Update weights for the next iteration
w_a, w_b, w_c = update_weights(w_a, w_b, w_c)

let w_a=w_a, w_b=w_b, w_c=w_c
    for i in 2:100
        mach.model.loss_function = LossContext(
            w_a=w_a, 
            w_b=w_b, 
            w_c=w_c
        )
        mach.model.niterations += 1
        fit!(mach)
        ...
        w_a, w_b, w_c = update_weights(w_a, w_b, w_c)
    end
end

And here are the warning messages I’m getting:

[ Info: Updating machine(SRRegressor(defaults = nothing, …), …).  
┌ Warning: Recreating population (output=1, population=2), as the saved one doesn't have the correct number of members.  
└ @ SymbolicRegression ~/.julia/packages/SymbolicRegression/kbyoS/src/SymbolicRegression.jl:743  
┌ Warning: Recreating population (output=1, population=16), as the saved one doesn't have the correct number of members.  
└ @ SymbolicRegression ~/.julia/packages/SymbolicRegression/kbyoS/src/SymbolicRegression.jl:743  
┌ Warning: Recreating population (output=1, population=18), as the saved one doesn't have the correct number of members.  
└ @ SymbolicRegression ~/.julia/packages/SymbolicRegression/kbyoS/src/SymbolicRegression.jl:743  
...

Would really appreciate any insights on whether this is expected behavior or if I should change something. Thanks again! 😊

[MilesCranmer]

Hm. I think it's fine. Maybe just increase by 2 (or more) iterations at a time instead of 1? Then it has a chance to repopulate before it needs to return the expressions

[GFODK]

Hey! It’s been about a week or two since I last asked, and I’ve been trying to mix in the ideas you suggested — thanks again for those tips!

I think I might’ve messed something up while doing that though. It seems like the weight updates inside my for loop (for i in 2:100) aren’t actually being applied properly. The loss function keeps using the initial values w_a = 1.0, w_b = 1.0, w_c = 1.0 instead of the updated ones.

Even though I’m assigning the updated weights like this:

for i in 2:100
    mach.model.loss_function = LossContext(
        w_a=w_a, 
        w_b=w_b, 
        w_c=w_c
    )
    ...

It looks like fit! isn’t picking up the updated loss_function values at all.

At first, I thought the issue might be with the let w_a=w_a, w_b=w_b, w_c=w_c block, so I tried removing that. But then I got this error:

ERROR: LoadError: UndefVarError: w_a not defined in local scope
Suggestion: check for an assignment to a local variable that shadows a global of the same name.

So I changed it to this:

w_a, w_b, w_c = update_weights(w_a, w_b, w_c)

for i in 2:100
    global w_a, w_b, w_c

    mach.model.loss_function = LossContext(
        w_a=w_a, 
        w_b=w_b, 
        w_c=w_c
    )
    mach.model.niterations += 1
    fit!(mach)
    ...

That removed the error message, but when I checked inside my custom loss function, the ctx.w_a, ctx.w_b, and ctx.w_c values were still stuck at 1.0 — so they weren’t updating after all.

Any idea why the updated weights aren’t getting passed into the loss function properly, even though I’m reassigning mach.model.loss_function each time? Let me know if I’m missing something obvious — I’ve been stuck on this one for a bit 😅

[MilesCranmer]

Can you share your full code? I can't reproduce based on this

[GFODK]

Thanks for the follow-up!
Actually, I should’ve just shared the full code earlier, but I didn’t include the full code earlier because I’m still working on writing things more efficiently, and the code got kind of long… I didn’t want to make things harder to follow with a huge block of messy code 😅

Anyway, here’s the full code I’m working with:
I’m trying to update the loss weights based on the minimum total loss among all the trees considered during a single iteration. So I structured the code this way to track that, and then update w_a, w_b, w_c accordingly.

Inside the for loop (starting from iteration 2), I’m reassigning loss_function = LossContext(...) using the weights calculated from the previous iteration’s loss. After training, I checked inside the custom loss function to see if the updated weights were actually used — but they still seemed stuck at the initial values. That’s why I’m suspecting the updated loss_function wasn’t properly applied.

Here’s the full code:

import SymbolicRegression: SRRegressor
import MLJ: machine, fit!, predict, report
using SymbolicRegression, DynamicDiff, DynamicExpressions, MLJ
using SymbolicRegression: Node
using DynamicDiff: D

import Zygote
using Base.Threads

Base.@kwdef struct LossContext <: Function
    w_a::Float64
    w_b::Float64
    w_c::Float64
end

function initialize_best_lists()
    global best_a_list  = Vector{Float64}()
    global best_b_list   = Vector{Float64}()
    global best_c_list  = Vector{Float64}()
    global best_w_a_list  = Vector{Float64}()
    global best_w_b_list  = Vector{Float64}()
    global best_w_c_list  = Vector{Float64}()

    global best_a_lock  = SpinLock()
    global best_b_lock  = SpinLock()
    global best_c_lock  = SpinLock()
    global best_w_a_lock  = SpinLock()
    global best_w_b_lock  = SpinLock()
    global best_w_c_lock  = SpinLock()
end

function save_list(lock_obj, data_list, value)
    lock(lock_obj) do
        push!(data_list, value)
    end
end

function update_weights(loss_lists::NTuple{4,Vector{Float64}}; beta_=1e-1)
    a_list, b_list, c_list = loss_lists

    function rate_of_change(lst)
        length(lst) < 2 ? 0.0 : lst[end] - lst[end-1]
    end
    
    function mean_loss(lst)
        isempty(lst) ? 0.0 : mean(lst)
    end

    s_i = [rate_of_change(a_list),
           rate_of_change(b_list),
           rate_of_change(c_list)]

    f_i = [mean_loss(a_list),
           mean_loss(b_list),
           mean_loss(c_list)]

    # Softmax-like
    exp_s_i = exp.(beta_ .* (s_i .- maximum(s_i)))
    w_i = exp_s_i ./ (sum(exp_s_i))

    weighted_w_i = f_i .* w_i
    sum_weighted_w_i = sum(weighted_w_i)
    w_i_ = weighted_w_i ./ sum_weighted_w_i

    return w_i_[1], w_i_[2], w_i_[3], w_i_[4]
end

    
function (ctx::LossContext)(tree::Node, dataset::Dataset{T,L}, options::Options) where {T,L}

    operators = OperatorEnum(binary_operators=(+, -, *, /), unary_operators=(sin, cos, exp, log))
    f = Expression(tree; operators=operators, variable_names=["x1","x2"])

    # Derivatives
    eq_dfdx   = D(f, 1)          # ∂f/∂x
    eq_d2fdx2 = D(eq_dfdx, 1)    # ∂²f/∂x²
    eq_dfdt   = D(eq_d2fdx2,

    f_val      = f(dataset.X)
    dfdx_val   = eq_dfdx(dataset.X)
    d2fdx2_val = eq_d2fdx2(dataset.X)

    # Reshape
    f_val      = reshape(f_val, size(dataset.X[1,:]))
    dfdx_val   = reshape(dfdx_val, size(dataset.X[1,:]))
    d2fdx2_val = reshape(d2fdx2_val, size(dataset.X[1,:]))

    # loss a
    L_a = sum((((2 .- dataset.X) .* d2fdx2_val) .- dfdx_val).^2) / dataset.n

    # loss b
    L_b = sum(f_val[begin].^2) / max(1, length(f_val[begin]))

    # loss c
    C = 1.0 # example constant
    L_c = sum(((dfdx_val[end]) .- C).^2) / max(1, length(dfdx_val[end]))

    # total loss
    loss = L_a * ctx.w_a + L_b * ctx.w_b + L_c * ctx.w_c

    # For checking the updates of each loss term and weight
    if !any(isnan, loss)
        if loss < last(best_loss_list)
            save_list(best_loss_lock, best_loss_list, loss)
            
            save_list(best_a_lock, best_a_list, a_loss)
            save_list(best_b_lock, best_b_list, b_loss)
            save_list(best_c_lock, best_c_list, c_loss)

            save_list(best_w_a_lock, best_w_a_list, ctx.w_a)
            save_list(best_w_b_lock, best_w_b_list, ctx.w_b)
            save_list(best_w_c_lock, best_w_c_list, ctx.w_c)
        end
    end

    return loss
end


w_a, w_b, w_c = 1.0, 1.0, 1.0
model = SRRegressor(
    niterations=1,
    binary_operators=[+, -, *, /],
    unary_operators=[sin, cos, exp, log],
    population_size=100,
    loss_function=LossContext(
        w_a=w_a, 
        w_b=w_b, 
        w_c=w_c
    )
)
mach = machine(model, x, y)
fit!(mach)


a_loss_list   = Vector{Float64}()
b_loss_list   = Vector{Float64}()
c_loss_list   = Vector{Float64}()
a_loss_lock   = SpinLock()
b_loss_lock   = SpinLock()
c_loss_lock   = SpinLock()

# To update the loss weight (only store the minimum total loss during each iteration)
save_list(a_loss_lock, a_loss_list, last(best_a_list))
save_list(b_loss_lock, b_loss_list, last(best_b_list))
save_list(c_loss_lock, c_loss_list, last(best_c_list))

loss_lists = (a_loss_list, b_loss_list, c_loss_list)

# Update weights for the next iteration
w_a, w_b, w_c = update_weights(loss_lists;)

let w_a=w_a, w_b=w_b, w_c=w_c
    for i in 2:100
        mach.model.loss_function = LossContext(
            w_a=w_a, 
            w_b=w_b, 
            w_c=w_c
        )
        mach.model.niterations += 1
        fit!(mach)

        save_list(a_loss_lock, a_loss_list, last(best_a_list))
        save_list(b_loss_lock, b_loss_list, last(best_b_list))
        save_list(c_loss_lock, c_loss_list, last(best_c_list))

        loss_lists = (a_loss_list, b_loss_list, c_loss_list)
        w_a, w_b, w_c = update_weights(loss_lists;)
    end
end

Let me know if anything looks off — I really appreciate the help!!

[GFODK]

Hi @MilesCranmer,
I went through the code line by line to figure out where the updated loss weight variables weren’t being passed correctly. Starting from the second iteration, I used a for loop, and I first confirmed that the weights were being updated with the following line:

w_a_updated, w_b_updated, w_c_updated = update_weights(losses_lists;)

After the loop for the second iteration finished, I checked the model using println(mach.model), and here’s what I saw:

SRRegressor(
  defaults = nothing, 
  binary_operators = Function[+, -, *, /], 
  unary_operators = Function[sin, cos, exp, log], 
  maxsize = 30, 
  maxdepth = nothing, 
  expression_type = Expression, 
  expression_options = NamedTuple(), 
  populations = nothing, 
  population_size = 100, 
  loss_function = LossContext(0.1660787615043684, 0.7782148007295373, 0.029241360342336993, 0.026463022423237673), 
  ...
  niterations = 2, 
  ...
)

From this, I could confirm that niterations was updated to 2, and the updated weights were also reflected in loss_function. However, when I printed out the weights inside the custom loss function like this:

println()
println()
println("A loss weight: ", ctx.w_a)
println("B loss weight: ", ctx.w_b)
println("C loss weight: ", ctx.w_c)
println()
println()

I still saw only the initial values of 1.0 being printed every time.

Now I’m wondering — could it be that internally mach.model is being cloned or copied, so the loss_function gets fixed at initialization and doesn’t actually update afterward? If that’s the case, and I need to recreate the machine object to apply the updated loss_function, wouldn’t that mean I’m losing all progress from the previous training iterations?

[MilesCranmer]

I am confused because obviously mach.model.niterations is changing. Argh. I wonder if it has to do with the model being stored in fitresult too? This might be a bug!

Can you try also doing

mach.model.loss_function = #= new loss =#
model.fitresult.model.loss_function = #= new loss =#

But in any case, there should be a proper way to update the model. I'm wondering if I introduced some bug with caching of the options or something.