Speed up HybridModel without changing the user API (#3)

CHANGELOG.md

@@ -1,6 +1,8 @@
 # Changelog
 
 ## Unreleased
+- Faster default `HybridModel` training without changing the user API. Mechanistic models stay keyword functions; the public `HybridModel` call still returns predictions, `parameters`, extras, and `nn_outputs`. On the synthetic RbQ10 example from the docs (neural `:rb`, global `:Q10`, hidden `[16, 16]`, sigmoid, `scale_nn_outputs=true`, `input_batchnorm=true`, batch 512) the training-loss gradient is about 3× faster than before (~1007 µs → ~315 µs) and about 1.17× a hand-written Lux model of the same math (was ~3×). `train` on 20 000 rows, 20 epochs, batch 512 goes from ~1.55 s to ~0.45 s. Only the default Vector `:mse` path (no `extra_loss`) is fused; the public forward is unchanged.
+- Default `train` keeps `NamedTuple` parameters from `Lux.setup`; a `ComponentArray` is used only for `AutoForwardDiff` and `Optimization.jl`. Whole-vector ops on `out.ps` (`length`, `norm`, `Array`) need `ComponentArray(ps)` first.
 
 ## v0.2.1 - 2026-09-04
 - Set Runic version to 1 in tools in order to match the one in CI. This should take care of the regular updates for patches, locally as well as remote. Older Manifest files in formatter should be updated always accordingly [#299](https://github.com/EarthyScience/EasyHybrid.jl/pull/299).

src/config/TrainingConfig.jl

@@ -22,8 +22,8 @@ $(TYPEDFIELDS)
     `Optimisers.jl`) three forms are accepted:
 
     1. A single `Optimisers.AbstractRule` (default), applied to the whole
-       parameter tree, e.g. `Adam(0.01)`. `ps` is wrapped in a
-       `ComponentArray`.
+       parameter tree, e.g. `Adam(0.01)`. `ps` is a `NamedTuple` unless
+       `autodiff_backend` is `AutoForwardDiff`.
     2. A `NamedTuple` of rules — one per top-level branch of the parameter
        tree — e.g. for an `RbQ10` hybrid model:
 

src/losses/compute_loss.jl

@@ -17,11 +17,23 @@ Main loss function for hybrid models that handles both training and evaluation m
 - In evaluation mode (`logging.train_mode = false`):
   - `(loss_values, st, ŷ)`: NamedTuple of losses, state and predictions
 """
+function compute_loss(HM::LuxCore.AbstractLuxContainerLayer, ps, st, data; logging::LoggingLoss)
+    return compute_loss(HM, ps, st, data, logging)
+end
+
+function compute_loss(
+        HM::HybridModel{<:Any, <:Vector}, ps, st,
+        ((x, forcings), (y_t, y_nan)),
+        logging::LoggingLoss{SymbolicLoss{S}, ExtraLoss{Nothing}, A, true}
+    ) where {S, A}
+    loss, st_nn = _hybrid_loss(HM, (x, forcings), ps, st, y_t, y_nan, Val(S), logging.agg)
+    return loss, ChainRulesCore.ignore_derivatives((; st_nn, fixed = st.fixed)), NamedTuple()
+end
+
 function compute_loss(
-        HM::LuxCore.AbstractLuxContainerLayer, ps, st, ((x, forcings), (y_t, y_nan));
+        HM::LuxCore.AbstractLuxContainerLayer, ps, st, ((x, forcings), (y_t, y_nan)),
         logging::LoggingLoss
     )
-
     targets = HM.targets
     ext_loss = extra_loss(logging)
     if logging.train_mode

src/losses/compute_loss_types.jl

@@ -3,9 +3,10 @@ export loss_types, training_loss, extra_loss
 
 abstract type LossSpec end
 
-struct SymbolicLoss <: LossSpec
+struct SymbolicLoss{S} <: LossSpec
     name::Symbol
 end
+SymbolicLoss(s::Symbol) = SymbolicLoss{s}(s)
 
 struct FunctionLoss <: LossSpec
     f::Function
@@ -21,8 +22,8 @@ ParameterizedLoss(f::Function) = ParameterizedLoss(f, (), NamedTuple())
 ParameterizedLoss(f::Function, args::Tuple) = ParameterizedLoss(f, args, NamedTuple())
 ParameterizedLoss(f::Function, kwargs::NamedTuple) = ParameterizedLoss(f, (), kwargs)
 
-struct ExtraLoss <: LossSpec
-    f::Union{Function, Nothing}
+struct ExtraLoss{F} <: LossSpec
+    f::F
 end
 
 """
@@ -105,10 +106,10 @@ logging = LoggingLoss(
 )
 ```
 """
-struct LoggingLoss{L <: Union{LossSpec, PerTarget}, T <: Function}
+struct LoggingLoss{L <: Union{LossSpec, PerTarget}, E <: LossSpec, T <: Function, TM}
     loss_types::Vector{LossSpec}
     training_loss::L
-    extra_loss::LossSpec
+    extra_loss::E
     agg::T
     train_mode::Bool
 end
@@ -125,11 +126,11 @@ function LoggingLoss(;
     tl = _to_loss_spec(training_loss)
     el = _to_extra_loss_spec(extra_loss)
 
-    return LoggingLoss{typeof(tl), F}(lt, tl, el, agg, train_mode)
+    return LoggingLoss{typeof(tl), typeof(el), F, train_mode}(lt, tl, el, agg, train_mode)
 end
 
 
-_to_loss_spec(s::Symbol) = SymbolicLoss(s)
+_to_loss_spec(s::Symbol) = SymbolicLoss{s}(s)
 _to_loss_spec(f::Function) = _accepts_params(f) ? ParamLoss(f) : FunctionLoss(f)
 _to_loss_spec(ls::LossSpec) = ls
 

src/models/GenericHybridModel.jl

@@ -33,7 +33,7 @@ It combines predictive neural networks (`NNs`) with a `mechanistic_model` to for
 
 $(TYPEDFIELDS)
 """
-struct HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)}
+struct HybridModel{T, P, MM, NP, GP, FP, KW, SN, TG} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)}
     "Neural network(s) used to predict parameters. Can be a single `Chain` or a `NamedTuple` of `Chain`s."
     NNs::T
 
@@ -47,7 +47,7 @@ struct HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)}
     targets::Vector{Symbol}
 
     "The core process-based or mechanistic model function."
-    mechanistic_model::Function
+    mechanistic_model::MM
 
     "Base parameters of the model (encapsulated in a `ParameterContainer`)."
     parameters::ParameterContainer
@@ -71,6 +71,26 @@ struct HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)}
     config::NamedTuple
 end
 
+function HybridModel(
+        NNs, predictors, forcing, targets, mechanistic_model, parameters,
+        neural_param_names, global_param_names, fixed_param_names,
+        scale_nn_outputs, start_from_default, config,
+    )
+    NP, GP, FP = Tuple(neural_param_names), Tuple(global_param_names), Tuple(fixed_param_names)
+    KW = _accepted_kwarg_names(
+        mechanistic_model,
+        Tuple(unique([forcing; neural_param_names; global_param_names; fixed_param_names])),
+    )
+    return HybridModel{
+        typeof(NNs), typeof(predictors), typeof(mechanistic_model),
+        NP, GP, FP, KW, scale_nn_outputs, Tuple(targets),
+    }(
+        NNs, predictors, forcing, targets, mechanistic_model, parameters,
+        neural_param_names, global_param_names, fixed_param_names,
+        scale_nn_outputs, start_from_default, config,
+    )
+end
+
 """
     _warn_forcing_param_overlap(forcing, param_names)
 
@@ -455,6 +475,75 @@ function _run_nn(m::HybridModel{<:Any, <:Vector}, ds_k::Tuple, ps, st)
     return scaled_nn_params, (; st_nn = st_nn), (;)
 end
 
+@generated function _hybrid_loss(
+        m::HybridModel{<:Any, <:Vector, MM, NP, GP, FP, KW, SN, TG},
+        ds_k::Tuple, ps, st, y, y_nan, ::Val{S}, agg
+    ) where {MM, NP, GP, FP, KW, SN, TG, S}
+    stmts = Expr[]
+    names = (NP..., GP..., FP...)
+    if NP === ()
+        push!(stmts, :(st_nn = st.st_nn))
+    else
+        push!(stmts, :((nn_out, st_nn) = LuxCore.apply(m.NNs, ds_k[1], ps.ps, st.st_nn)))
+        push!(stmts, :(slices = eachslice(nn_out, dims = 1)))
+    end
+    for (i, n) in enumerate(NP)
+        rhs = SN ?
+            :(scale_single_param(Val($(QuoteNode(n))), slices[$i], m.parameters)) :
+            :(slices[$i])
+        push!(stmts, Expr(:(=), n, rhs))
+    end
+    for n in GP
+        push!(stmts, Expr(:(=), n, :(scale_single_param(Val($(QuoteNode(n))), getproperty(ps, $(QuoteNode(n))), m.parameters))))
+    end
+    for n in FP
+        push!(stmts, Expr(:(=), n, :(getproperty(st.fixed, $(QuoteNode(n))))))
+    end
+    NP === () && GP === () && push!(stmts, :(ChainRulesCore.ignore_derivatives(ps)))
+    NP === () && KW === () && push!(stmts, :(ChainRulesCore.ignore_derivatives(ds_k)))
+    !SN && GP === () && push!(stmts, :(ChainRulesCore.ignore_derivatives(m.parameters)))
+    if KW === nothing
+        push!(stmts, :(all_params = (; $(names...))))
+        push!(stmts, :(y_pred = m.mechanistic_model(; merge(ds_k[2], all_params)...)))
+    else
+        args = Expr[]
+        for n in KW
+            if n in names
+                push!(args, Expr(:kw, n, n))
+            else
+                push!(args, Expr(:kw, n, :(getfield(ds_k[2], $(QuoteNode(n))))))
+            end
+        end
+        push!(stmts, :(y_pred = m.mechanistic_model(; $(args...))))
+    end
+    if TG === ()
+        push!(stmts, :(ChainRulesCore.ignore_derivatives((y, y_nan, agg))))
+        push!(stmts, :(return agg(()), st_nn))
+    else
+        terms = Expr[]
+        for t in TG
+            qt = QuoteNode(t)
+            yi = y <: NamedTuple ? :(getfield(y, $qt)) : :(_get_target_y(y, $qt))
+            nani = y_nan <: NamedTuple ? :(getfield(y_nan, $qt)) : :(_get_target_y(y_nan, $qt))
+            loss = S === :mse ?
+                :(mean(abs2, ŷ_i[nan_i] .- y_i[nan_i])) :
+                :(_apply_loss(ŷ_i, y_i, nan_i, $(QuoteNode(S))))
+            push!(
+                terms, :(
+                    let
+                        y_i = $yi
+                        ŷ_i = _get_target_ŷ(y_pred, y_i, $qt)
+                        nan_i = $nani
+                        $loss
+                    end
+                )
+            )
+        end
+        push!(stmts, :(return agg(($(terms...),)), st_nn))
+    end
+    return Expr(:block, stmts...)
+end
+
 """
     (m::HybridModel)(ds_k::Tuple, ps, st)
 
@@ -510,7 +599,7 @@ function (m::HybridModel)(ds_k::Tuple, ps, st)
     out = (; y_pred..., extra_params..., parameters = all_params, out_extra...)
     st_new = (; st_new_nns..., fixed = st.fixed)
 
-    return out, st_new
+    return out, ChainRulesCore.ignore_derivatives(st_new)
 end
 
 function (m::HybridModel)(ds_k, ps, st)

src/models/utils.jl

@@ -76,10 +76,18 @@ scaletype(hm::ParameterContainer, name) = hm.scales[name]
 Map an unconstrained `raw_val` into the parameter's bounds `[ℓ, u]`, using the
 warp selected for `name` (`:linear`, `:log`, or `:logit`).
 """
-function scale_single_param(name, raw_val, hm::ParameterContainer)
-    return _scale(Val(scaletype(hm, name)), lower(hm)[name], upper(hm)[name], raw_val)
+function scale_single_param(::Val{N}, raw_val, hm::ParameterContainer) where {N}
+    bounds = getfield(hm.values, N)
+    ℓ = bounds[2]
+    u = bounds[3]
+    return ℓ .+ (u .- ℓ) .* sigmoid.(raw_val)
 end
 
+scale_single_param(name::Symbol, raw_val, hm::ParameterContainer) =
+    scale_single_param(Val(name), raw_val, hm)
+
+inv_sigmoid(y) = log.(y ./ (1 .- y))
+
 """ 
     scale_single_param_minmax(name, hm::ParameterContainer)
 

src/training/epoch.jl

@@ -51,7 +51,7 @@ function build_loss_fn(model, cfg::TrainConfig)
         extra_loss = cfg.extra_loss,
         agg = cfg.agg
     )
-    return (model, ps, st, (x, y)) -> compute_loss(model, ps, st, (x, y); logging = logging)
+    return (model, ps, st, (x, y)) -> compute_loss(model, ps, st, (x, y), logging)
 end
 
 function evaluate_epoch(model, x_train, forcings_train, y_train, mask_train, x_val, forcings_val, y_val, mask_val, ps, st, epoch::Int, init::EpochSnapshot, cfg::TrainConfig)

src/training/initialization.jl

@@ -44,7 +44,7 @@ function init_model_state(model, cfg::TrainConfig)
             model, ps, st, cfg.opt, opt_state, 0,
         )
     elseif is_optimisers_rule(cfg.opt)
-        ps = ps |> ComponentArray
+        cfg.autodiff_backend isa Lux.AutoForwardDiff && (ps = ps |> ComponentArray)
         Lux.Training.TrainState(model, ps, st, cfg.opt)
     else
         ps = ps |> ComponentArray

src/training/train_optimization.jl

@@ -127,7 +127,7 @@ function _build_optim_loss(model, st, cfg::TrainConfig)
         agg = cfg.agg,
     )
     return function (p, data)
-        loss_val, _, _ = compute_loss(model, p, st, data; logging)
+        loss_val, _, _ = compute_loss(model, p, st, data, logging)
         return loss_val
     end
 end

src/utils/ad_rules.jl

@@ -0,0 +1,6 @@
+# Zygote traces Base's keyword-arg name sorting; the results are Bool/Symbol.
+for (fname, nargs) in ((:diff_names, 2), (:sym_in, 2), (:merge_names, 2))
+    isdefined(Base, fname) || continue
+    args = ntuple(_ -> :(::Any), nargs)
+    @eval ChainRulesCore.@non_differentiable Base.$fname($(args...))
+end

src/utils/utils.jl

@@ -1,3 +1,4 @@
+include("ad_rules.jl")
 include("macro_hybrid.jl")
 include("wrap_tuples.jl")
 include("tools.jl")

test/runtests.jl

@@ -17,4 +17,5 @@ include("test_wrap_tuples.jl")
 include("test_extract_weights.jl")
 include("test_transformers.jl")
 include("test_seq2seq.jl")
+include("test_synth_performance.jl")
 include("test_zarr_io.jl")

test/test_compute_loss.jl

@@ -5,6 +5,7 @@ using Statistics
 using DimensionalData
 using Random
 using DataFrames
+using Zygote
 
 @testset "_compute_loss" begin
     # Test data setup
@@ -305,6 +306,16 @@ end
             ŷ_actual, y_t, y_nan, targets, :mse, sum
         )
         @test loss_value ≈ main_loss
+
+        g_fused = only(Zygote.gradient(p -> compute_loss(HM, p, st, (data[1], (data[2], y_nan)), logging)[1], ps))
+        g_pub = only(
+            Zygote.gradient(
+                p -> _compute_loss(HM(data[1], p, st)[1], y_t, y_nan, Tuple(targets), :mse, sum),
+                ps,
+            )
+        )
+        @test g_fused.b ≈ g_pub.b
+        @test g_fused.ps.layer_2.weight ≈ g_pub.ps.layer_2.weight
     end
 
     @testset "Evaluation mode with extra_loss" begin

test/test_generic_hybrid_model.jl

@@ -2,6 +2,7 @@ using Lux
 using Random
 using AxisKeys
 using ComponentArrays
+using Zygote
 
 # Test data generation
 dk = gen_linear_data()
@@ -255,6 +256,11 @@ end
         @test haskey(output.parameters, :c)
         @test haskey(new_st, :st_nn)
         @test haskey(new_st, :fixed)
+
+        # Gradient of a prediction-only loss (returned state unused)
+        g = Zygote.gradient(p -> sum(abs2, model(data_[1], p, st)[1].y_pred), ps)
+        @test !isnothing(g[1])
+        @test g[1].b isa AbstractVector
     end
 
     @testset "HybridModel with scale_nn_outputs=true" begin