From e3439f4820dd29ef4951cfc8e971697c7ec70c37 Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Tue, 5 May 2026 14:05:20 +0200
Subject: [PATCH 01/16] support kwargs and configs

---
 src/training/train.jl | 89 +++++++++++++++++++++++++------------------
 1 file changed, 52 insertions(+), 37 deletions(-)

diff --git a/src/training/train.jl b/src/training/train.jl
index a16abc25..b2e8fe0f 100644
--- a/src/training/train.jl
+++ b/src/training/train.jl
@@ -1,41 +1,6 @@
 export train
 
-"""
-    train(model, data; train_cfg::TrainConfig = TrainConfig(), data_cfg::DataConfig = DataConfig())
-
-Train a hybrid model using the provided data.
-
-Returns `nothing` if data preparation fails (zero-size dimension in training or validation data).
-
-# Arguments
-- `model`: The hybrid model to train.
-- `data`: Training data, a single `DimArray`, a single `DataFrame`, or a single `KeyedArray`.
-
-# Keyword Arguments
-- `train_cfg`: Training configuration. See [`TrainConfig`](@ref) for all options.
-- `data_cfg`: Data preparation configuration. See [`DataConfig`](@ref) for all options.
-
-# Returns
-A [`TrainResults`](@ref) with the following fields:
-- `train_losses`: Per-epoch training losses.
-- `val_losses`: Per-epoch validation losses.
-- `snapshots`: Model parameter snapshots taken during training.
-- `train_obs_pred`: Observed vs. predicted values on the training set.
-- `val_obs_pred`: Observed vs. predicted values on the validation set.
-- `train_diffs`: Additional diagnostic variables computed on the training set.
-- `val_diffs`: Additional diagnostic variables computed on the validation set.
-- `ps`: Final (or best) model parameters.
-- `st`: Final (or best) model state.
-- `best_epoch`: Epoch at which the best validation loss was achieved.
-- `best_loss`: Best validation loss recorded during training.
-
-# Example
-```julia
-cfg = TrainConfig(nepochs=100, batchsize=32)
-result = train(myModel, myData; train_cfg=cfg)
-```
-"""
-function train(model, data; train_cfg::TrainConfig = TrainConfig(), data_cfg::DataConfig = DataConfig())
+function _train(model, data, train_cfg::TrainConfig, data_cfg::DataConfig)
     validate_config(train_cfg)
     ext = load_makie_extension(train_cfg)
     seed!(train_cfg.random_seed)
@@ -78,6 +43,56 @@ function train(model, data; train_cfg::TrainConfig = TrainConfig(), data_cfg::Da
     return build_results(model, history, stopper, ps, st, x_train, forcings_train, y_train, x_val, forcings_val, y_val, train_cfg)
 end
 
+"""
+    train(model, data; train_cfg::TrainConfig = TrainConfig(), data_cfg::DataConfig = DataConfig())
+    train(model, data; kwargs...)   # kwargs forwarded to `TrainConfig` / `DataConfig`
+
+Train a hybrid model using the provided data.
+
+Returns `nothing` if data preparation fails (zero-size dimension in training or validation data).
+
+# Arguments
+- `model`: The hybrid model to train.
+- `data`: Training data, a single `DimArray`, a single `DataFrame`, or a single `KeyedArray`.
+
+# Keyword Arguments
+- `train_cfg`: Training configuration. See [`TrainConfig`](@ref) for all options.
+- `data_cfg`: Data preparation configuration. See [`DataConfig`](@ref) for all options.
+- Any other kwargs (deprecated) are forwarded as fields to `TrainConfig` / `DataConfig`.
+
+# Returns
+A [`TrainResults`](@ref) with the following fields:
+- `train_losses`: Per-epoch training losses.
+- `val_losses`: Per-epoch validation losses.
+- `snapshots`: Model parameter snapshots taken during training.
+- `train_obs_pred`: Observed vs. predicted values on the training set.
+- `val_obs_pred`: Observed vs. predicted values on the validation set.
+- `train_diffs`: Additional diagnostic variables computed on the training set.
+- `val_diffs`: Additional diagnostic variables computed on the validation set.
+- `ps`: Final (or best) model parameters.
+- `st`: Final (or best) model state.
+- `best_epoch`: Epoch at which the best validation loss was achieved.
+- `best_loss`: Best validation loss recorded during training.
+
+# Example
+```julia
+cfg = TrainConfig(nepochs=100, batchsize=32)
+result = train(myModel, myData; train_cfg=cfg)
+```
+"""
+function train(
+    model, data;
+    train_cfg::TrainConfig = TrainConfig(),
+    data_cfg::DataConfig  = DataConfig(),
+    kwargs...,
+)
+    if !isempty(kwargs)
+        train_cfg, data_cfg = kwargs_to_configs((), kwargs)
+    end
+
+    return _train(model, data, train_cfg, data_cfg)
+end
+
 function valid_mask(y)
     nt = (;)
     isempty = true
@@ -108,7 +123,7 @@ function train(model, data, save_ps; kwargs...)
     )
 
     train_cfg, data_cfg = kwargs_to_configs(save_ps, kwargs)
-    return train(model, data; train_cfg, data_cfg)
+    return _train(model, data, train_cfg, data_cfg)
 end
 
 function expand_sequence_kwargs(kwargs)

From 4af107bfb47bf29216b4964cf106ba5f933f53b2 Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Tue, 5 May 2026 14:32:08 +0200
Subject: [PATCH 02/16] kwargs over_ride configs

---
 src/training/train.jl | 100 ++++++++++++++++++++++++++++++------------
 1 file changed, 73 insertions(+), 27 deletions(-)

diff --git a/src/training/train.jl b/src/training/train.jl
index b2e8fe0f..c6fd90e9 100644
--- a/src/training/train.jl
+++ b/src/training/train.jl
@@ -45,10 +45,31 @@ end
 
 """
     train(model, data; train_cfg::TrainConfig = TrainConfig(), data_cfg::DataConfig = DataConfig())
-    train(model, data; kwargs...)   # kwargs forwarded to `TrainConfig` / `DataConfig`
+    train(model, data; kwargs...)
 
 Train a hybrid model using the provided data.
 
+Two equivalent calling styles are supported:
+
+1. **Typed configs** — pass complete `TrainConfig` / `DataConfig` objects:
+   ```julia
+   train(model, data;
+       train_cfg = TrainConfig(nepochs=100, batchsize=32),
+       data_cfg  = DataConfig(split_data_at=0.8),
+   )
+   ```
+
+2. **Flat kwargs** — pass `TrainConfig` / `DataConfig` field names directly:
+   ```julia
+   train(model, data; nepochs=100, batchsize=32, split_data_at=0.8)
+   ```
+
+The two styles can also be mixed; flat kwargs override the corresponding fields of
+the supplied `train_cfg` / `data_cfg`:
+```julia
+train(model, data; train_cfg = TrainConfig(nepochs=100), nepochs = 10)  # nepochs = 10
+```
+
 Returns `nothing` if data preparation fails (zero-size dimension in training or validation data).
 
 # Arguments
@@ -58,7 +79,7 @@ Returns `nothing` if data preparation fails (zero-size dimension in training or
 # Keyword Arguments
 - `train_cfg`: Training configuration. See [`TrainConfig`](@ref) for all options.
 - `data_cfg`: Data preparation configuration. See [`DataConfig`](@ref) for all options.
-- Any other kwargs (deprecated) are forwarded as fields to `TrainConfig` / `DataConfig`.
+- Any other kwargs are forwarded as overrides to `train_cfg` / `data_cfg`.
 
 # Returns
 A [`TrainResults`](@ref) with the following fields:
@@ -73,12 +94,6 @@ A [`TrainResults`](@ref) with the following fields:
 - `st`: Final (or best) model state.
 - `best_epoch`: Epoch at which the best validation loss was achieved.
 - `best_loss`: Best validation loss recorded during training.
-
-# Example
-```julia
-cfg = TrainConfig(nepochs=100, batchsize=32)
-result = train(myModel, myData; train_cfg=cfg)
-```
 """
 function train(
     model, data;
@@ -87,9 +102,8 @@ function train(
     kwargs...,
 )
     if !isempty(kwargs)
-        train_cfg, data_cfg = kwargs_to_configs((), kwargs)
+        train_cfg, data_cfg = override_configs(train_cfg, data_cfg, kwargs)
     end
-
     return _train(model, data, train_cfg, data_cfg)
 end
 
@@ -107,20 +121,6 @@ function valid_mask(y)
 end
 
 function train(model, data, save_ps; kwargs...)
-    Base.depwarn(
-        """
-        `train(model, data, save_ps; kwargs...)` is deprecated.
-        Use the new API instead:
-
-            train(model, data;
-                train_cfg = TrainConfig(nepochs=100, ...),
-                data_cfg  = DataConfig(split_data_at=0.8, ...)
-            )
-
-        See `?TrainConfig` and `?DataConfig` for all available options.
-        """,
-        :train
-    )
 
     train_cfg, data_cfg = kwargs_to_configs(save_ps, kwargs)
     return _train(model, data, train_cfg, data_cfg)
@@ -130,7 +130,6 @@ function expand_sequence_kwargs(kwargs)
     haskey(kwargs, :sequence_kwargs) || return kwargs
 
     seq_kw = kwargs[:sequence_kwargs]
-    @warn "`sequence_kwargs` is deprecated, pass sequence options directly via `DataConfig` instead."
 
     # map old sequence_kwargs keys to new DataConfig field names
     key_map = Dict(
@@ -144,12 +143,18 @@ function expand_sequence_kwargs(kwargs)
         get(key_map, k, k) => v for (k, v) in pairs(seq_kw)
     )
 
-    # drop sequence_kwargs, merge expanded fields
     remaining = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k !== :sequence_kwargs)
     return merge(remaining, expanded)
 end
 
-function kwargs_to_configs(save_ps, kwargs)
+"""
+    kwargs_to_configs(kwargs) -> (TrainConfig, DataConfig)
+
+Build a fresh `(TrainConfig, DataConfig)` pair from a flat collection of kwargs.
+Kwargs are split between the two configs based on `fieldnames(TrainConfig)` and
+`fieldnames(DataConfig)`; unknown kwargs are warned about and dropped.
+"""
+function kwargs_to_configs(kwargs)
     train_keys = fieldnames(TrainConfig)
     data_keys = fieldnames(DataConfig)
 
@@ -172,6 +177,47 @@ function kwargs_to_configs(save_ps, kwargs)
     return TrainConfig(; train_kwargs...), DataConfig(; data_kwargs...)
 end
 
+"""
+    override_configs(train_cfg, data_cfg, kwargs)
+
+Return `(train_cfg′, data_cfg′)` where any field present in `kwargs` overrides the
+corresponding field of `train_cfg`/`data_cfg`. Fields not mentioned in `kwargs` are
+kept as-is. Unknown kwargs trigger a warning and are ignored.
+
+This is the merge step used by the deprecated kwargs path of `train` — it lets users
+mix the new typed API (`train_cfg=...`) with leftover individual kwargs, with kwargs
+taking precedence.
+"""
+function override_configs(train_cfg::TrainConfig, data_cfg::DataConfig, kwargs)
+    train_keys = fieldnames(TrainConfig)
+    data_keys = fieldnames(DataConfig)
+
+    kwargs = rename_deprecated_kwargs(kwargs)
+    kwargs = expand_sequence_kwargs(kwargs)
+
+    train_overrides = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k in train_keys)
+    data_overrides  = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k in data_keys)
+
+    unknown = [k for k in keys(kwargs) if k ∉ train_keys && k ∉ data_keys]
+    if !isempty(unknown)
+        @warn "Unknown kwargs will be ignored: $(join(unknown, ", "))"
+    end
+
+    return override_config(train_cfg, train_overrides), override_config(data_cfg, data_overrides)
+end
+
+"""
+    override_config(cfg, overrides::NamedTuple)
+
+Return a new `cfg` of the same type with the fields named in `overrides` replaced.
+Works with any `@kwdef` struct (e.g. [`TrainConfig`](@ref), [`DataConfig`](@ref)).
+"""
+function override_config(cfg::T, overrides::NamedTuple) where {T}
+    isempty(overrides) && return cfg
+    base = NamedTuple(k => getfield(cfg, k) for k in fieldnames(T))
+    return T(; merge(base, overrides)...)
+end
+
 const DEPRECATED_KWARG_NAMES = Dict(
     :file_name => :model_name,
     :hybrid_name => :model_name,

From 714d6ea539747ba12f33b78f7576ce44425d62fd Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Fri, 8 May 2026 11:40:32 +0200
Subject: [PATCH 03/16] Refactor save_ps_st function to include epoch parameter
 for saving model state to not get epoch already exists error

---
 src/io/checkpoints.jl | 2 +-
 src/io/save.jl        | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/io/checkpoints.jl b/src/io/checkpoints.jl
index 9131c4e8..7b7555de 100644
--- a/src/io/checkpoints.jl
+++ b/src/io/checkpoints.jl
@@ -20,7 +20,7 @@ function save_final!(paths::TrainingPaths, model, ps, st, x_train, forcings_trai
     if cfg.save_training
         target_names = model.targets
         save_epoch = stopper.best_epoch == 0 ? 0 : stopper.best_epoch
-        save_ps_st!(paths.best_model, model, cfg.cdev(ps), cfg.cdev(st), cfg.tracked_params, save_epoch)
+        save_ps_st(paths.best_model, model, cfg.cdev(ps), cfg.cdev(st), cfg.tracked_params, save_epoch)
 
         ŷ_train, αst_train = model((cfg.cdev(x_train), cfg.cdev(forcings_train)), cfg.cdev(ps), LuxCore.testmode(cfg.cdev(st)))
         ŷ_val, αst_val = model((cfg.cdev(x_val), cfg.cdev(forcings_val)), cfg.cdev(ps), LuxCore.testmode(cfg.cdev(st)))
diff --git a/src/io/save.jl b/src/io/save.jl
index 325c82f0..f5c37133 100644
--- a/src/io/save.jl
+++ b/src/io/save.jl
@@ -1,6 +1,6 @@
 export get_all_groups
 export load_group
-function save_ps_st(file_name, hm, ps, st, save_ps)
+function save_ps_st(file_name, hm, ps, st, save_ps, epoch = 0)
     hm_name = string(nameof(typeof(hm)))
     # split physical parameters
     tmp_e = if !isempty(save_ps)
@@ -10,9 +10,9 @@ function save_ps_st(file_name, hm, ps, st, save_ps)
 
     isfile(file_name) && rm(file_name)
     return jldopen(file_name, "w") do file
-        file["HybridModel_$hm_name/epoch_0"] = (ps, st)
+        file["HybridModel_$hm_name/epoch_$epoch"] = (ps, st)
         if !isempty(save_ps)
-            file["physical_params/epoch_0"] = tmp_e
+            file["physical_params/epoch_$epoch"] = tmp_e
         end
     end
 end

From 0495dd1fb933eecbf3c537c3959467e531255647 Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Fri, 8 May 2026 12:00:35 +0200
Subject: [PATCH 04/16] kwargs to config two pos args

---
 src/training/train.jl | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/training/train.jl b/src/training/train.jl
index c6fd90e9..1ec5e1b7 100644
--- a/src/training/train.jl
+++ b/src/training/train.jl
@@ -148,13 +148,16 @@ function expand_sequence_kwargs(kwargs)
 end
 
 """
-    kwargs_to_configs(kwargs) -> (TrainConfig, DataConfig)
+    kwargs_to_configs(save_ps, kwargs) -> (TrainConfig, DataConfig)
 
 Build a fresh `(TrainConfig, DataConfig)` pair from a flat collection of kwargs.
 Kwargs are split between the two configs based on `fieldnames(TrainConfig)` and
 `fieldnames(DataConfig)`; unknown kwargs are warned about and dropped.
+
+`save_ps` is the deprecated positional argument from `train(model, data, save_ps; ...)`;
+when non-empty it is forwarded as `tracked_params` on the resulting `TrainConfig`.
 """
-function kwargs_to_configs(kwargs)
+function kwargs_to_configs(save_ps, kwargs)
     train_keys = fieldnames(TrainConfig)
     data_keys = fieldnames(DataConfig)
 

From a777da7e7f374739e48c765da35989a4a4a0e116 Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Fri, 8 May 2026 14:52:08 +0200
Subject: [PATCH 05/16] parameter bounds in yaml, mechanistic model in yaml

---
 src/config/config_yaml.jl | 101 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 99 insertions(+), 2 deletions(-)

diff --git a/src/config/config_yaml.jl b/src/config/config_yaml.jl
index fbf4136e..9bb78643 100644
--- a/src/config/config_yaml.jl
+++ b/src/config/config_yaml.jl
@@ -1,5 +1,5 @@
 export load_hybrid_config, save_hybrid_config
-export get_hybrid_config, get_train_config, get_full_config
+export get_hybrid_config, get_train_config, get_full_config, get_parameters_config, get_mechanistic_model_config
 
 function load_hybrid_config(path::String; dicttype = OrderedDict{String, Any})
     return load_file(path; dicttype)
@@ -12,11 +12,108 @@ end
 function get_hybrid_config(hm::LuxCore.AbstractLuxContainerLayer)
     hm_config = OrderedDict{String, Any}()
     for field in fieldnames(typeof(hm))
-        hm_config[string(field)] = getfield(hm, field)
+        hm_config[string(field)] = _yaml_field_value(getfield(hm, field))
     end
     return hm_config
 end
 
+_yaml_field_value(x) = x
+_yaml_field_value(p::AbstractHybridModel) = get_parameters_config(p)
+_yaml_field_value(f::Function) = get_mechanistic_model_config(f)
+
+"""
+    get_parameters_config(p::AbstractHybridModel)
+
+Serialize the parameter table (`default`, `lower`, `upper` per parameter) of a
+`HybridParams`/`ParameterContainer` into a nested `OrderedDict` suitable for
+YAML output. Without this, only the compact `show` of the struct
+(e.g. `HybridParams(ParameterContainer(RUE, Rb, Q10))`) would be written, which
+drops all of the actual default and bound values.
+"""
+function get_parameters_config(p::AbstractHybridModel)
+    pc = hasfield(typeof(p), :hybrid) ? p.hybrid : p
+    out = OrderedDict{String, Any}()
+    for name in keys(pc.values)
+        d, l, u = pc.values[name]
+        out[string(name)] = OrderedDict{String, Any}(
+            "default" => d,
+            "lower" => l,
+            "upper" => u,
+        )
+    end
+    return out
+end
+
+"""
+    get_mechanistic_model_config(f::Function)
+
+Build an `OrderedDict` describing a function for YAML output: its `name`, and
+for each `Method`, its source `file`, starting `line`, and the full `source`
+text extracted from disk by parsing one complete expression starting at that
+line. The single-method case is flattened so the YAML stays compact.
+
+Used to record `mechanistic_model` in the saved config so the exact function
+definition (not just its name) is preserved alongside the run.
+"""
+function get_mechanistic_model_config(f::Function)
+    out = OrderedDict{String, Any}()
+    out["name"] = string(nameof(f))
+    method_entries = OrderedDict{String, Any}[]
+    for m in methods(f)
+        entry = OrderedDict{String, Any}()
+        try
+            file, line = Base.functionloc(m)
+            entry["file"] = string(file)
+            entry["line"] = line
+            src = _try_extract_function_source(string(file), line)
+            if src !== nothing
+                entry["source"] = src
+            end
+        catch err
+            entry["error"] = sprint(showerror, err)
+        end
+        push!(method_entries, entry)
+    end
+    if length(method_entries) == 1
+        merge!(out, method_entries[1])
+    elseif !isempty(method_entries)
+        out["methods"] = method_entries
+    end
+    return out
+end
+
+# Read `file` and return the source text of the first complete top-level
+# expression starting at `line`. Returns `nothing` if the file can't be read or
+# no expression can be parsed (e.g. function defined at the REPL or inside an
+# eval'd string).
+function _try_extract_function_source(file::AbstractString, line::Integer)
+    isfile(file) || return nothing
+    text = try
+        read(file, String)
+    catch
+        return nothing
+    end
+    idx = firstindex(text)
+    cur = 1
+    while cur < line
+        nl = findnext(==('\n'), text, idx)
+        nl === nothing && return nothing
+        idx = nextind(text, nl)
+        cur += 1
+    end
+    expr_and_next = try
+        Meta.parse(text, idx; greedy = true, raise = false)
+    catch
+        return nothing
+    end
+    expr_and_next === nothing && return nothing
+    _, next_idx = expr_and_next
+    last_idx = lastindex(text)
+    endidx = next_idx > last_idx ? last_idx : prevind(text, next_idx)
+    endidx < idx && return nothing
+    return rstrip(text[idx:endidx])
+end
+
 function get_train_config(train_args::NamedTuple)
     train_config = OrderedDict{String, Any}()
     for field in fieldnames(typeof(train_args))

From 15ed5758ca9cd7190b6d7b797594a9698cb8f178 Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Mon, 18 May 2026 11:42:54 +0200
Subject: [PATCH 06/16] towards regularization

---
 src/utils/extract_weights.jl | 37 ++++++++++++++++++++++++++++++++++++
 src/utils/utils.jl           |  1 +
 2 files changed, 38 insertions(+)
 create mode 100644 src/utils/extract_weights.jl

diff --git a/src/utils/extract_weights.jl b/src/utils/extract_weights.jl
new file mode 100644
index 00000000..35fc7963
--- /dev/null
+++ b/src/utils/extract_weights.jl
@@ -0,0 +1,37 @@
+export extract_weights
+using ComponentArrays
+
+"""
+    extract_weights(ps; key=:weight) -> Vector{AbstractArray}
+
+Walk the parameter tree `ps` (a `ComponentArray`, `NamedTuple`, or any nested
+combination of them) and return all leaf arrays whose **immediate parent field
+name** equals `key`.
+
+Defaults to `:weight`, so you get the weight matrices of `Dense`/`Conv` layers
+and skip biases, `BatchNorm` `scale`/`bias`, running statistics in `st`, and
+any scalar global parameters.
+
+The returned arrays are views/aliases into `ps`. When `ps` is the argument the
+autodiff is differentiating w.r.t., gradients of any function of these views
+flow back into the trainable weights.
+"""
+function extract_weights(ps; key::Symbol = :weight)
+    out = AbstractArray[]
+    _collect!(out, ps, key)
+    return out
+end
+
+_collect!(_, _, ::Symbol) = nothing
+
+function _collect!(out, node::Union{NamedTuple, ComponentArray}, key::Symbol)
+    for name in propertynames(node)
+        child = getproperty(node, name)
+        if name === key && child isa AbstractArray && isempty(propertynames(child))
+            push!(out, child)
+        else
+            _collect!(out, child, key)
+        end
+    end
+    return nothing
+end
\ No newline at end of file
diff --git a/src/utils/utils.jl b/src/utils/utils.jl
index 1d38e2d1..1389bb97 100644
--- a/src/utils/utils.jl
+++ b/src/utils/utils.jl
@@ -5,3 +5,4 @@ include("plotrecipes.jl")
 include("helpers_data_loading.jl")
 include("helpers_cross_validation.jl")
 include("print_banner.jl")
+include("extract_weights.jl")
\ No newline at end of file

From 4feb1cdd3e390e1ffda526946bf5cc5a127657ca Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Tue, 19 May 2026 10:45:54 +0200
Subject: [PATCH 07/16] L2 loss steps

---
 .../research/synthetic_respiration.jl         |  4 +-
 src/config/TrainingConfig.jl                  |  2 +-
 src/losses/compute_loss.jl                    |  4 +-
 src/losses/compute_loss_types.jl              |  2 +-
 src/utils/extract_weights.jl                  | 50 +++++++++++++++++--
 test/runtests.jl                              |  1 +
 test/test_compute_loss.jl                     | 32 ++++++++++--
 test/test_extract_weights.jl                  | 33 ++++++++++++
 8 files changed, 114 insertions(+), 14 deletions(-)
 create mode 100644 test/test_extract_weights.jl

diff --git a/docs/literate/research/synthetic_respiration.jl b/docs/literate/research/synthetic_respiration.jl
index df6c0f4a..a9c57651 100644
--- a/docs/literate/research/synthetic_respiration.jl
+++ b/docs/literate/research/synthetic_respiration.jl
@@ -89,7 +89,7 @@ single_nn_hybrid_model = constructHybridModel(
 
 # ### train on DataFrame
 
-extra_loss = function (ŷ)
+function rb_prior_extra_loss(ŷ, ps)
     return (; a = sum((5.0 .- ŷ.rb) .^ 2))
 end
 
@@ -106,7 +106,7 @@ single_nn_out = train(
     shuffleobs = true,
     loss_types = [:mse, :nse],
     show_progress = false,
-    extra_loss = extra_loss,
+    extra_loss = rb_prior_extra_loss,
     model_name = "RbQ10_synthetic1"
 )
 
diff --git a/src/config/TrainingConfig.jl b/src/config/TrainingConfig.jl
index e3a2f6f9..6ce9dc33 100644
--- a/src/config/TrainingConfig.jl
+++ b/src/config/TrainingConfig.jl
@@ -44,7 +44,7 @@ loss computation, data handling, output, and visualization.
     """
     loss_types::Vector{Symbol} = [:mse, :r2]
 
-    "Additional loss term to add to the training loss. Default: `nothing`."
+    "Additional loss `(ŷ, ps; kwargs...) -> NamedTuple` added to the training loss. Default: `nothing`."
     extra_loss = nothing
 
     "Aggregation function applied to computed losses. Default: `sum`."
diff --git a/src/losses/compute_loss.jl b/src/losses/compute_loss.jl
index 580391d8..74aa59da 100644
--- a/src/losses/compute_loss.jl
+++ b/src/losses/compute_loss.jl
@@ -29,7 +29,7 @@ function compute_loss(
         loss_value = _compute_loss(ŷ, y_t, y_nan, targets, training_loss(logging), logging.agg)
         # Add extra_loss if provided
         if ext_loss !== nothing
-            extra_loss_value = ext_loss(ŷ)
+            extra_loss_value = ext_loss(ŷ, ps)
             loss_value = logging.agg([loss_value, extra_loss_value...])
         end
         stats = NamedTuple()
@@ -38,7 +38,7 @@ function compute_loss(
         loss_value = _compute_loss(ŷ, y_t, y_nan, targets, loss_types(logging), logging.agg)
         # Add extra_loss entries if provided
         if ext_loss !== nothing
-            extra_loss_values = ext_loss(ŷ)
+            extra_loss_values = ext_loss(ŷ, ps)
             agg_extra_loss_value = logging.agg(extra_loss_values)
             loss_value = (; loss_value..., extra_loss = (; extra_loss_values..., Symbol(logging.agg) => agg_extra_loss_value))
         end
diff --git a/src/losses/compute_loss_types.jl b/src/losses/compute_loss_types.jl
index 0b476148..8245d56e 100644
--- a/src/losses/compute_loss_types.jl
+++ b/src/losses/compute_loss_types.jl
@@ -58,7 +58,7 @@ A structure to define a logging loss function for hybrid models.
     - `(f, kwargs)`: keyword args, e.g. `(scaled_loss, (scale=2.0,))`
     - `(f, args, kwargs)`: both, e.g. `(complex_loss, (0.5,), (scale=2.0,))`
 - `training_loss`: The loss specification to use during training (same format as above)
-- `extra_loss`: Optional function `(ŷ; kwargs...) -> scalar` to add to training loss (default: `nothing`)
+- `extra_loss`: Optional function `(ŷ, ps; kwargs...) -> NamedTuple` (or splattable collection) added to training loss (default: `nothing`)
 - `agg`: Function to aggregate losses across targets, e.g. `sum` or `mean`
 - `train_mode`: If true, uses `training_loss`; otherwise uses `loss_types`.
 
diff --git a/src/utils/extract_weights.jl b/src/utils/extract_weights.jl
index 35fc7963..91eeb501 100644
--- a/src/utils/extract_weights.jl
+++ b/src/utils/extract_weights.jl
@@ -1,6 +1,16 @@
-export extract_weights
+export extract_weights, weight_l2
 using ComponentArrays
 
+"""
+    _is_named_leaf(name, child, key) -> Bool
+
+Return `true` when `child` is an array leaf whose parent field is `key`
+(e.g. a `Dense` layer's `weight` matrix).
+"""
+function _is_named_leaf(name::Symbol, child, key::Symbol)
+    return name === key && child isa AbstractArray && isempty(propertynames(child))
+end
+
 """
     extract_weights(ps; key=:weight) -> Vector{AbstractArray}
 
@@ -27,11 +37,45 @@ _collect!(_, _, ::Symbol) = nothing
 function _collect!(out, node::Union{NamedTuple, ComponentArray}, key::Symbol)
     for name in propertynames(node)
         child = getproperty(node, name)
-        if name === key && child isa AbstractArray && isempty(propertynames(child))
+        if _is_named_leaf(name, child, key)
             push!(out, child)
         else
             _collect!(out, child, key)
         end
     end
     return nothing
-end
\ No newline at end of file
+end
+
+"""
+    weight_l2(ps; key=:weight) -> Real
+
+Sum of squared Frobenius norms over all parameter arrays in `ps` whose
+immediate parent field name is `key` (default `:weight`).
+
+Unlike `sum(abs2, extract_weights(ps))`, this fuses the tree walk with the
+reduction so it is safe to use inside Zygote-differentiated losses, e.g.:
+
+```julia
+extra_loss = (ŷ, ps) -> (; l2_Rb = λ * weight_l2(ps.Rb),)
+```
+
+When `ps` is the loss function argument, gradients flow into the weight arrays.
+"""
+function weight_l2(ps; key::Symbol = :weight)
+    return _weight_l2(ps, key)
+end
+
+_weight_l2(::Any, ::Symbol) = 0f0
+
+function _weight_l2(node::Union{NamedTuple, ComponentArray}, key::Symbol)
+    s = 0f0
+    for name in propertynames(node)
+        child = getproperty(node, name)
+        if _is_named_leaf(name, child, key)
+            s = s + sum(abs2, child)
+        else
+            s = s + _weight_l2(child, key)
+        end
+    end
+    return s
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 8ba2c9fd..210088f2 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -13,6 +13,7 @@ include("test_loss_fn.jl")
 include("test_show_train.jl")
 include("test_show_generic_hybrid.jl")
 include("test_wrap_tuples.jl")
+include("test_extract_weights.jl")
 
 @testset "LinearHM" begin
     # test model instantiation
diff --git a/test/test_compute_loss.jl b/test/test_compute_loss.jl
index b7cb2915..81a7f64f 100644
--- a/test/test_compute_loss.jl
+++ b/test/test_compute_loss.jl
@@ -255,8 +255,9 @@ end
     y_nan = (; var1 = trues(n_samples), var2 = trues(n_samples))
 
     @testset "Training mode with extra_loss" begin
-        # Define extra loss function
-        extra_loss_func(ŷ) = [sum(abs, ŷ.var1), sum(abs, ŷ.var2)]
+        function extra_loss_func(ŷ, ps)
+            return [sum(abs, ŷ.var1), sum(abs, ŷ.var2)]
+        end
 
         logging = LoggingLoss(
             loss_types = [:mse],
@@ -278,7 +279,7 @@ end
         main_loss = _compute_loss(
             ŷ_actual, y_t, y_nan, targets, :mse, sum
         )
-        extra_loss_vals = extra_loss_func(ŷ_actual)
+        extra_loss_vals = extra_loss_func(ŷ_actual, ps)
         expected_loss = sum([main_loss, extra_loss_vals...])
         @test loss_value ≈ expected_loss
     end
@@ -307,8 +308,9 @@ end
     end
 
     @testset "Evaluation mode with extra_loss" begin
-        # Define extra loss function that returns a NamedTuple
-        extra_loss_func(ŷ) = (var1_extra = sum(abs, ŷ.var1), var2_extra = sum(abs, ŷ.var2))
+        function extra_loss_func(ŷ, ps)
+            return (var1_extra = sum(abs, ŷ.var1), var2_extra = sum(abs, ŷ.var2))
+        end
 
         logging = LoggingLoss(
             loss_types = [:mse, :mae],
@@ -358,4 +360,24 @@ end
         @test haskey(stats, :var1)
         @test haskey(stats, :var2)
     end
+
+    @testset "Training mode with weight_l2 extra_loss" begin
+        λ = 1.0f-4
+        function nn_weight_extra_loss(ŷ, ps)
+            return (; l2_ps = λ * weight_l2(ps.ps))
+        end
+
+        logging = LoggingLoss(
+            loss_types = [:mse],
+            training_loss = :mse,
+            extra_loss = nn_weight_extra_loss,
+            train_mode = true
+        )
+
+        loss_value, _, _ = compute_loss(HM, ps, st, (data[1], (data[2], y_nan)); logging = logging)
+        ŷ_actual, _ = HM(data[1], ps, st)
+        main_loss = _compute_loss(ŷ_actual, y_t, y_nan, targets, :mse, sum)
+        expected = main_loss + λ * weight_l2(ps.ps)
+        @test loss_value ≈ expected
+    end
 end
diff --git a/test/test_extract_weights.jl b/test/test_extract_weights.jl
new file mode 100644
index 00000000..6029f75f
--- /dev/null
+++ b/test/test_extract_weights.jl
@@ -0,0 +1,33 @@
+using EasyHybrid
+using Lux
+using Random
+using Test
+using Zygote
+
+@testset "extract_weights / weight_l2" begin
+    rng = Random.default_rng()
+    chain = Chain(
+        BatchNorm(2),
+        Dense(2 => 16, tanh),
+        Dense(16 => 1),
+    )
+    ps, _ = Lux.setup(rng, chain)
+
+    ws = extract_weights(ps)
+    @test !isempty(ws)
+    @test all(w -> w isa AbstractMatrix, ws)
+
+    manual_l2 = sum(sum(abs2, w) for w in ws)
+    @test weight_l2(ps) ≈ manual_l2
+
+    λ = 1.0f-3
+    loss(ps) = λ * weight_l2(ps)
+    g = Zygote.gradient(loss, ps)[1]
+    @test g !== nothing
+    g_ws = extract_weights(g)
+    @test !isempty(g_ws)
+    @test all(gw -> all(!iszero, gw), g_ws)
+
+    # biases are not regularized by default
+    @test weight_l2(ps; key = :bias) ≈ sum(sum(abs2, b) for b in extract_weights(ps; key = :bias))
+end

From 359479349c1ae1b4e5f1e138390c69cfe728aac8 Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Wed, 20 May 2026 15:23:22 +0200
Subject: [PATCH 08/16] LSTM fix time is not sample anymore

---
 src/data/split_data.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/data/split_data.jl b/src/data/split_data.jl
index 330097bc..b561f74b 100644
--- a/src/data/split_data.jl
+++ b/src/data/split_data.jl
@@ -168,6 +168,8 @@ function collect_end_dim(x_all::Union{KeyedArray{Float32, 3}, AbstractDimArray{F
     return collect(getindex(x_all, :, :, idx))
 end
 
+# 2D (feature, time): split along time; 3D (feature, time, batch): split along batch
+_num_samples(x::AbstractArray{<:Any, 3}) = size(x, 3)
 _num_samples(x::AbstractArray) = size(x, 2)
 _num_samples(x::NamedTuple) = _num_samples(first(values(x)))
 

From bdb8e54073bc6c8354d8869ec840579d613f88fa Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Tue, 26 May 2026 10:42:24 +0200
Subject: [PATCH 09/16] per parameter optimiser

---
 src/config/TrainingConfig.jl   |  50 +++++++++-
 src/training/initialization.jl |  28 +++++-
 src/training/train.jl          | 162 +++++++++++++++++++++++++++------
 3 files changed, 207 insertions(+), 33 deletions(-)

diff --git a/src/config/TrainingConfig.jl b/src/config/TrainingConfig.jl
index 6ce9dc33..82a9360e 100644
--- a/src/config/TrainingConfig.jl
+++ b/src/config/TrainingConfig.jl
@@ -13,7 +13,33 @@ loss computation, data handling, output, and visualization.
     "Size of the training batches. Default: 64."
     batchsize::Int = 64
 
-    "Optimizer to use for training. Default: Adam(0.01)."
+    """
+    Optimizer to use for training. Default: `Adam(0.01)`.
+
+    On the `Optimisers.jl` path (selected when `opt` originates from
+    `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`.
+    2. A `NamedTuple` of rules — one per top-level branch of the parameter
+       tree — e.g. for an `RbQ10` hybrid model:
+
+       ```julia
+       opt = (; Rb = Adam(1e-3), Q10 = Descent(1e-2))
+       ```
+
+       The framework calls `Optimisers.setup` per branch; branches not
+       listed fall back to `Adam()`. `ps` is kept as a nested `NamedTuple`
+       (no `ComponentArray` wrap) because `Optimisers.jl` treats a
+       `ComponentArray` as a single leaf and would collapse the per-branch
+       state tree.
+    3. A `NamedTuple` of pre-built state trees (already returned by
+       `Optimisers.setup`). Useful when one branch needs an
+       `Optimisers.OptimiserChain` or a frozen leaf built up by hand.
+
+    Forms 2 and 3 can be freely mixed in the same `NamedTuple`.
+    """
     opt = Adam(0.01)
 
     """
@@ -97,6 +123,28 @@ loss computation, data handling, output, and visualization.
 
     "Tuple of parameter names to track across epochs. Default: `()`."
     tracked_params::Tuple = ()
+
+    """
+    `Optimization.jl` path only: pass the full training set as a single tuple
+    to `OptimizationProblem` instead of an `MLUtils.DataLoader`. Has no effect
+    on the `Optimisers.jl` path. Default: `false`.
+    """
+    full_batch::Bool = false
+
+    """
+    `Optimization.jl` path only: promote `ps` to `Float64` before optimization.
+    Mirrors the workaround documented in `projects/RbQ10/Q10_lbfgs.jl` for
+    [Lux.jl#1260](https://github.com/LuxDL/Lux.jl/issues/1260). Has no effect
+    on the `Optimisers.jl` path. Default: `false`.
+    """
+    promote_f64::Bool = false
+
+    """
+    `Optimization.jl` path only: build a validation `EpochSnapshot` (driving
+    history / early-stopping / dashboard) every `eval_every` callback hits.
+    Default: `1`.
+    """
+    eval_every::Int = 1
 end
 
 function validate_config(cfg::TrainConfig)
diff --git a/src/training/initialization.jl b/src/training/initialization.jl
index 6d8c0c6d..d2feb6b4 100644
--- a/src/training/initialization.jl
+++ b/src/training/initialization.jl
@@ -17,13 +17,35 @@ end
 function init_model_state(model, cfg::TrainConfig)
     if isnothing(cfg.train_from)
         ps, st = LuxCore.setup(Random.default_rng(), model)
-        ps = ps |> ComponentArray
     else
         ps, st = get_ps_st(cfg.train_from)
     end
 
-    # ps = ComponentArray(ps)
-    train_state = Lux.Training.TrainState(model, ps, st, cfg.opt)
+    train_state = if is_per_branch_opt(cfg.opt)
+        # Keep `ps` as a nested `NamedTuple`: `Optimisers.jl` treats a
+        # `ComponentArray` as a single leaf (no `Functors.functor` method),
+        # which would collapse the per-branch state tree into one big Leaf
+        # and silently apply the first rule to everything.
+        ps isa ComponentArray && (ps = NamedTuple(ps))
+        opt_state = build_opt_state(cfg.opt, ps)
+        # `Lux.Training.TrainState`'s public constructor accepts only
+        # `Optimisers.AbstractRule`. We use the (internal but stable)
+        # positional constructor of `@concrete struct TrainState` to inject
+        # a pre-built opt_state. Field order follows Lux 1.x; see
+        # `Lux.Training.TrainState` in `Lux/src/helpers/training.jl`.
+        Lux.Training.TrainState(
+            nothing,    # cache
+            nothing,    # objective_function
+            nothing,    # allocator_cache
+            model, ps, st, cfg.opt, opt_state, 0,
+        )
+    elseif is_optimisers_rule(cfg.opt)
+        ps = ps |> ComponentArray
+        Lux.Training.TrainState(model, ps, st, cfg.opt)
+    else
+        ps = ps |> ComponentArray
+        nothing
+    end
 
     return ps, st, train_state
 end
diff --git a/src/training/train.jl b/src/training/train.jl
index 1ec5e1b7..62840bfd 100644
--- a/src/training/train.jl
+++ b/src/training/train.jl
@@ -1,5 +1,97 @@
 export train
 
+"""
+    is_optimisers_rule(opt) -> Bool
+
+Return `true` when `opt` originates from the `Optimisers.jl` package (e.g.
+`Adam`, `AdamW`, `RMSProp`, `OptimiserChain`), **or** when `opt` is a
+`NamedTuple` describing a per-branch optimizer (see [`is_per_branch_opt`](@ref)).
+
+The check on single rules is by source package
+(`nameof(parentmodule(typeof(opt))) === :Optimisers`) rather than by
+`isa Optimisers.AbstractRule`, because in some package combinations
+`Optim.jl` optimizers were observed to satisfy the `AbstractRule` test and
+get misrouted to the `Lux.Training` loop.
+
+The `Lux.Training`-based loop dispatches on this; everything else (including
+`Optim.jl` and `Optimization.jl` optimizers) is routed through the
+`Optimization.jl` driver.
+"""
+is_optimisers_rule(opt) =
+    nameof(parentmodule(typeof(opt))) === :Optimisers ||
+    is_per_branch_opt(opt)
+
+"""
+    is_per_branch_opt(opt) -> Bool
+
+Return `true` when `opt` is a `NamedTuple` describing a per-branch
+optimizer specification — i.e. one of:
+
+  - a `NamedTuple` of `Optimisers.AbstractRule`s (e.g.
+    `(; Rb = Adam(1e-3), Q10 = Descent(1e-2))`), or
+  - a `NamedTuple` of pre-built optimizer state trees as returned by
+    `Optimisers.setup(rule, ps_branch)`, or
+  - a mix of the two.
+
+Branches of the parameter tree not listed in `opt` fall back to the default
+rule `Adam()` (see [`build_opt_state`](@ref)).
+
+This is detected purely by `opt isa NamedTuple`; the per-branch dispatch is
+checked again per-leaf when [`build_opt_state`](@ref) walks the spec.
+"""
+is_per_branch_opt(opt) = opt isa NamedTuple
+
+"""
+    build_opt_state(opt, ps::NamedTuple; default_rule = Optimisers.Adam())
+
+Build the optimizer state tree consumed by `Lux.Training.TrainState` /
+`Optimisers.update!`. Three forms of `opt` are accepted:
+
+  1. `opt::Optimisers.AbstractRule` — single rule applied to the whole
+     parameter tree (delegates to `Optimisers.setup(opt, ps)`).
+  2. `opt::NamedTuple` of `Optimisers.AbstractRule`s — each rule is wired
+     to the matching top-level branch of `ps` via
+     `Optimisers.setup(rule, ps[name])`. Branches missing from `opt` use
+     `default_rule`.
+  3. `opt::NamedTuple` of pre-built state trees (already returned by a
+     prior `Optimisers.setup`) — used as-is. Form 2 and 3 can be mixed in
+     the same `NamedTuple`.
+
+The returned state tree has the same top-level keys as `ps`.
+
+# Example
+
+```julia
+ps, _ = LuxCore.setup(rng, hybrid_model)   # (; Rb = NN_ps, RUE = NN_ps, Q10 = [v])
+
+# Form 2 — preferred, lets the framework call `Optimisers.setup`:
+opt_state = build_opt_state(
+    (; Rb = Adam(1e-3), RUE = Adam(1e-3), Q10 = Descent(1e-2)),
+    ps,
+)
+```
+"""
+function build_opt_state(opt::Optimisers.AbstractRule, ps; default_rule = Optimisers.Adam())
+    return Optimisers.setup(opt, ps)
+end
+
+function build_opt_state(opt::NamedTuple, ps::NamedTuple; default_rule = Optimisers.Adam())
+    pairs_ = map(collect(keys(ps))) do k
+        if haskey(opt, k)
+            spec = opt[k]
+            return k => spec isa Optimisers.AbstractRule ?
+                Optimisers.setup(spec, getproperty(ps, k)) :
+                spec
+        else
+            return k => Optimisers.setup(default_rule, getproperty(ps, k))
+        end
+    end
+    extra = setdiff(collect(keys(opt)), collect(keys(ps)))
+    isempty(extra) ||
+        @warn "Per-branch optimizer keys not found in parameter tree, ignored: $(extra)"
+    return NamedTuple(pairs_)
+end
+
 function _train(model, data, train_cfg::TrainConfig, data_cfg::DataConfig)
     validate_config(train_cfg)
     ext = load_makie_extension(train_cfg)
@@ -43,6 +135,27 @@ function _train(model, data, train_cfg::TrainConfig, data_cfg::DataConfig)
     return build_results(model, history, stopper, ps, st, x_train, forcings_train, y_train, x_val, forcings_val, y_val, train_cfg)
 end
 
+"""
+    _train(model, data, train_cfg, data_cfg, solve_kwargs)
+
+Dispatcher used by `train(...)`: routes to the original 4-arg `_train` body
+(the `Lux.Training` / `Optimisers.jl` loop) when `train_cfg.opt isa
+Optimisers.AbstractRule`, or to `_train_optimization` (which delegates batch
+iteration to `Optimization.jl`) otherwise. `solve_kwargs` are forwarded to
+`solve(...)` on the `Optimization.jl` branch and warned about on the
+`Optimisers.jl` branch.
+"""
+function _train(model, data, train_cfg::TrainConfig, data_cfg::DataConfig, solve_kwargs::NamedTuple)
+    return if is_optimisers_rule(train_cfg.opt)
+        if !isempty(solve_kwargs)
+            @warn "Unknown kwargs ignored on the Optimisers.jl path: $(join(keys(solve_kwargs), ", "))"
+        end
+        _train(model, data, train_cfg, data_cfg)
+    else
+        _train_optimization(model, data, train_cfg, data_cfg, solve_kwargs)
+    end
+end
+
 """
     train(model, data; train_cfg::TrainConfig = TrainConfig(), data_cfg::DataConfig = DataConfig())
     train(model, data; kwargs...)
@@ -101,10 +214,8 @@ function train(
     data_cfg::DataConfig  = DataConfig(),
     kwargs...,
 )
-    if !isempty(kwargs)
-        train_cfg, data_cfg = override_configs(train_cfg, data_cfg, kwargs)
-    end
-    return _train(model, data, train_cfg, data_cfg)
+    train_cfg, data_cfg, solve_kwargs = override_configs(train_cfg, data_cfg, kwargs)
+    return _train(model, data, train_cfg, data_cfg, solve_kwargs)
 end
 
 function valid_mask(y)
@@ -121,9 +232,8 @@ function valid_mask(y)
 end
 
 function train(model, data, save_ps; kwargs...)
-
-    train_cfg, data_cfg = kwargs_to_configs(save_ps, kwargs)
-    return _train(model, data, train_cfg, data_cfg)
+    train_cfg, data_cfg, solve_kwargs = kwargs_to_configs(save_ps, kwargs)
+    return _train(model, data, train_cfg, data_cfg, solve_kwargs)
 end
 
 function expand_sequence_kwargs(kwargs)
@@ -148,11 +258,13 @@ function expand_sequence_kwargs(kwargs)
 end
 
 """
-    kwargs_to_configs(save_ps, kwargs) -> (TrainConfig, DataConfig)
+    kwargs_to_configs(save_ps, kwargs) -> (TrainConfig, DataConfig, NamedTuple)
 
 Build a fresh `(TrainConfig, DataConfig)` pair from a flat collection of kwargs.
 Kwargs are split between the two configs based on `fieldnames(TrainConfig)` and
-`fieldnames(DataConfig)`; unknown kwargs are warned about and dropped.
+`fieldnames(DataConfig)`; anything left over is returned as the third element
+and forwarded to `solve(...)` on the `Optimization.jl` path (or warned about
+on the `Optimisers.jl` path — see `_train`).
 
 `save_ps` is the deprecated positional argument from `train(model, data, save_ps; ...)`;
 when non-empty it is forwarded as `tracked_params` on the resulting `TrainConfig`.
@@ -166,30 +278,24 @@ function kwargs_to_configs(save_ps, kwargs)
 
     train_kwargs = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k in train_keys)
     data_kwargs = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k in data_keys)
-
-    unknown = [k for k in keys(kwargs) if k ∉ train_keys && k ∉ data_keys]
-    if !isempty(unknown)
-        @warn "Unknown kwargs will be ignored: $(join(unknown, ", "))"
-    end
+    solve_kwargs = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k ∉ train_keys && k ∉ data_keys)
 
     if !isempty(save_ps)
         @warn "`save_ps` is deprecated, use `TrainConfig(tracked_params=(...))` instead."
         train_kwargs = merge(train_kwargs, (; tracked_params = save_ps))
     end
 
-    return TrainConfig(; train_kwargs...), DataConfig(; data_kwargs...)
+    return TrainConfig(; train_kwargs...), DataConfig(; data_kwargs...), solve_kwargs
 end
 
 """
-    override_configs(train_cfg, data_cfg, kwargs)
-
-Return `(train_cfg′, data_cfg′)` where any field present in `kwargs` overrides the
-corresponding field of `train_cfg`/`data_cfg`. Fields not mentioned in `kwargs` are
-kept as-is. Unknown kwargs trigger a warning and are ignored.
+    override_configs(train_cfg, data_cfg, kwargs) -> (TrainConfig, DataConfig, NamedTuple)
 
-This is the merge step used by the deprecated kwargs path of `train` — it lets users
-mix the new typed API (`train_cfg=...`) with leftover individual kwargs, with kwargs
-taking precedence.
+Return `(train_cfg′, data_cfg′, solve_kwargs)` where any field present in `kwargs`
+overrides the corresponding field of `train_cfg`/`data_cfg`. Fields not mentioned
+in `kwargs` are kept as-is. Anything left over is returned as `solve_kwargs` and
+forwarded to `solve(...)` on the `Optimization.jl` path (or warned about on the
+`Optimisers.jl` path — see `_train`).
 """
 function override_configs(train_cfg::TrainConfig, data_cfg::DataConfig, kwargs)
     train_keys = fieldnames(TrainConfig)
@@ -200,13 +306,11 @@ function override_configs(train_cfg::TrainConfig, data_cfg::DataConfig, kwargs)
 
     train_overrides = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k in train_keys)
     data_overrides  = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k in data_keys)
+    solve_kwargs    = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k ∉ train_keys && k ∉ data_keys)
 
-    unknown = [k for k in keys(kwargs) if k ∉ train_keys && k ∉ data_keys]
-    if !isempty(unknown)
-        @warn "Unknown kwargs will be ignored: $(join(unknown, ", "))"
-    end
-
-    return override_config(train_cfg, train_overrides), override_config(data_cfg, data_overrides)
+    return override_config(train_cfg, train_overrides),
+           override_config(data_cfg, data_overrides),
+           solve_kwargs
 end
 
 """

From 19142cf54347de2a523b758aa6b92a6cb6bede93 Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Mon, 1 Jun 2026 14:42:45 +0200
Subject: [PATCH 10/16] hard sigmoid inverse

---
 src/models/GenericHybridModel.jl | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl
index 1bd093b7..20bb5e2b 100644
--- a/src/models/GenericHybridModel.jl
+++ b/src/models/GenericHybridModel.jl
@@ -1,4 +1,4 @@
-export SingleNNHybridModel, MultiNNHybridModel, constructHybridModel, scale_single_param, AbstractHybridModel, build_hybrid, ParameterContainer, default, lower, upper, hard_sigmoid, inv_sigmoid
+export SingleNNHybridModel, MultiNNHybridModel, constructHybridModel, scale_single_param, AbstractHybridModel, build_hybrid, ParameterContainer, default, lower, upper, hard_sigmoid, inv_hard_sigmoid, inv_sigmoid
 export HybridParams
 
 # Import necessary components for neural networks
@@ -10,6 +10,13 @@ function hard_sigmoid(x)
     return clamp.(0.2 .* x .+ 0.5, 0.0, 1.0)
 end
 
+# Inverse of `hard_sigmoid` on the linear region (0, 1).
+# Saturated inputs (y ≤ 0 or y ≥ 1) are extrapolated linearly since the
+# clamp makes the forward map non-invertible there.
+function inv_hard_sigmoid(y)
+    return (y .- 0.5) ./ 0.2
+end
+
 abstract type AbstractHybridModel end
 
 mutable struct ParameterContainer{NT <: NamedTuple, T} <: AbstractHybridModel

From b2daf950d2c0cc223aea4770662531d219863014 Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Tue, 2 Jun 2026 10:46:01 +0200
Subject: [PATCH 11/16] multi lstm

---
 src/data/sequences.jl            | 69 ++++++++++++++++++++++++++++++++
 src/models/GenericHybridModel.jl |  5 ++-
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/src/data/sequences.jl b/src/data/sequences.jl
index 64e42ee1..ba51c7aa 100644
--- a/src/data/sequences.jl
+++ b/src/data/sequences.jl
@@ -34,6 +34,37 @@ end
 _subset_y(y, valid) = y[:, :, valid]
 _subset_y(y::NamedTuple, valid) = map(v -> v[:, valid], y)
 
+# NamedTuple predictors (MultiNNHybridModel): every branch is a 3D
+# `(feature, time, batch)` array sharing the batch axis with forcings/targets.
+# Drop a sequence if any branch has a NaN predictor or all targets are NaN.
+function filter_sequences(x_tuple::Tuple{<:NamedTuple, <:Any}, y)
+    x, forcings = x_tuple
+    valid = _valid_seq_indices(x, y)
+    nseq = _nseq(x)
+    length(valid) < nseq && @info "Dropped $(nseq - length(valid)) / $nseq sequences with NaN predictors or all-NaN targets"
+    x_filtered = map(b -> b[:, :, valid], x)
+    forcings_filtered = map(v -> v[:, valid], forcings)
+    return (x_filtered, forcings_filtered), _subset_y(y, valid)
+end
+
+_nseq(x::NamedTuple) = size(first(values(x)), 3)
+
+_branches_finite(x::NamedTuple, ii) = all(b -> !any(isnan, @view(b[:, :, ii])), values(x))
+_targets_present(y, ii) = any(!isnan, @view(y[:, :, ii]))
+_targets_present(y::NamedTuple, ii) = any(v -> any(!isnan, @view(v[:, ii])), values(y))
+
+# `(x::NamedTuple, y::NamedTuple)` is defined explicitly to disambiguate against
+# `_valid_seq_indices(x, y::NamedTuple)` above.
+function _valid_seq_indices(x::NamedTuple, y::NamedTuple)
+    n = _nseq(x)
+    return findall(ii -> _branches_finite(x, ii) && _targets_present(y, ii), 1:n)
+end
+
+function _valid_seq_indices(x::NamedTuple, y)
+    n = _nseq(x)
+    return findall(ii -> _branches_finite(x, ii) && _targets_present(y, ii), 1:n)
+end
+
 """
     split_into_sequences(x, y; input_window=5, output_window=1, output_shift=1, lead_time=1)
 
@@ -68,6 +99,44 @@ function split_into_sequences(x::Tuple, y; kwargs...)
     return (X_seq, forcings_seq), Y_seq
 end
 
+# MultiNNHybridModel support.
+#
+# For `MultiNNHybridModel`, `prepare_data` returns the predictors as a
+# NamedTuple of per-branch `(feature, time)` arrays (one per neural network),
+# while forcings and targets stay shared across branches. Window every branch on
+# the common time axis and keep the NamedTuple structure so the model's forward
+# pass can index `ds_k[1][nn_name]`.
+function split_into_sequences(x::Tuple{<:NamedTuple, <:Any}, y::NamedTuple; kwargs...)
+    x_branches, forcings = x
+    isempty(x_branches) && throw(ArgumentError("predictor NamedTuple is empty"))
+    targetkeys = keys(y)
+
+    # All branches share the same time axis; use the first as the reference for
+    # windowing the (shared) targets and forcings.
+    ref = _as_keyed_time(first(values(x_branches)))
+    y_arr = _nt_to_array(y, ref)
+
+    X_seq = map(b -> first(split_into_sequences(_as_keyed_time(b), y_arr; kwargs...)), x_branches)
+    _, Y_seq = split_into_sequences(ref, y_arr; kwargs...)
+    forcings_seq = _window_forcings(forcings, ref; kwargs...)
+
+    return (X_seq, forcings_seq), _array_to_nt(Y_seq, targetkeys)
+end
+
+# Predictor branches built for `MultiNNHybridModel` are plain matrices (their
+# axis keys are stripped in `prepare_data`). Wrap them with synthetic keys so
+# they flow through the keyed windowing path and produce keyed 3D outputs that
+# `collect_end_dim`/`view_end_dim` understand. Keyed/Dim inputs pass through.
+function _as_keyed_time(b)
+    (b isa KeyedArray || b isa AbstractDimArray) && return b
+    ndims(b) == 2 || throw(ArgumentError("expected a 2D (feature, time) predictor branch; got ndims = $(ndims(b))"))
+    return KeyedArray(
+        Float32.(b);
+        variable = Symbol.("f", 1:size(b, 1)),
+        time = 1:size(b, 2),
+    )
+end
+
 function split_into_sequences(x, y::NamedTuple; kwargs...)
     targetkeys = keys(y)
     y_arr = _nt_to_array(y, x)
diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl
index 20bb5e2b..5bf56e0c 100644
--- a/src/models/GenericHybridModel.jl
+++ b/src/models/GenericHybridModel.jl
@@ -483,7 +483,10 @@ function (m::MultiNNHybridModel)(ds_k::Tuple, ps, st)
     # 4) Scale neural network parameters using the mapping
     scaled_nn_params = NamedTuple()
     for (nn_name, param_name) in zip(keys(m.NNs), m.neural_param_names)
-        nn_cols = eachrow(nn_outputs[nn_name])
+        # `eachslice(...; dims = 1)` (instead of `eachrow`) so this works for both
+        # the feed-forward case (2D `(param, batch)` output) and the recurrent/LSTM
+        # case (3D `(param, time, batch)` sequence output).
+        nn_cols = eachslice(nn_outputs[nn_name]; dims = 1)
 
         # Create parameter for this NN
         nn_param = NamedTuple{(param_name,), Tuple{typeof(nn_cols[1])}}((nn_cols[1],))

From 40690f9e66d8dd655b2700d4ee899107978174c7 Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Wed, 10 Jun 2026 10:28:24 +0200
Subject: [PATCH 12/16] l2

---
 src/utils/extract_weights.jl | 28 +++++++++++++++++++---------
 test/test_extract_weights.jl |  9 +++++++++
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/src/utils/extract_weights.jl b/src/utils/extract_weights.jl
index 91eeb501..78de9727 100644
--- a/src/utils/extract_weights.jl
+++ b/src/utils/extract_weights.jl
@@ -8,7 +8,9 @@ Return `true` when `child` is an array leaf whose parent field is `key`
 (e.g. a `Dense` layer's `weight` matrix).
 """
 function _is_named_leaf(name::Symbol, child, key::Symbol)
-    return name === key && child isa AbstractArray && isempty(propertynames(child))
+    # ComponentArray views (SubArray, ReshapedArray, …) have non-empty
+    # `propertynames`; only recurse into nested ComponentArray/NamedTuple nodes.
+    return name === key && child isa AbstractArray && !(child isa ComponentArray)
 end
 
 """
@@ -47,35 +49,43 @@ function _collect!(out, node::Union{NamedTuple, ComponentArray}, key::Symbol)
 end
 
 """
-    weight_l2(ps; key=:weight) -> Real
+    weight_l2(ps; key=:weight, normalize=false) -> Real
 
 Sum of squared Frobenius norms over all parameter arrays in `ps` whose
 immediate parent field name is `key` (default `:weight`).
 
+With `normalize=true`, returns the mean squared weight (sum divided by the
+number of scalar weights), so the value is independent of network width/depth.
+
 Unlike `sum(abs2, extract_weights(ps))`, this fuses the tree walk with the
 reduction so it is safe to use inside Zygote-differentiated losses, e.g.:
 
 ```julia
-extra_loss = (ŷ, ps) -> (; l2_Rb = λ * weight_l2(ps.Rb),)
+extra_loss = (ŷ, ps) -> (; l2_Rb = λ * weight_l2(ps.Rb; normalize=true),)
 ```
 
 When `ps` is the loss function argument, gradients flow into the weight arrays.
 """
-function weight_l2(ps; key::Symbol = :weight)
-    return _weight_l2(ps, key)
+function weight_l2(ps; key::Symbol = :weight, normalize::Bool = false)
+    s, n = _weight_l2_stats(ps, key)
+    return normalize ? (n > 0 ? s / n : zero(s)) : s
 end
 
-_weight_l2(::Any, ::Symbol) = 0f0
+_weight_l2_stats(::Any, ::Symbol) = (0f0, 0)
 
-function _weight_l2(node::Union{NamedTuple, ComponentArray}, key::Symbol)
+function _weight_l2_stats(node::Union{NamedTuple, ComponentArray}, key::Symbol)
     s = 0f0
+    n = 0
     for name in propertynames(node)
         child = getproperty(node, name)
         if _is_named_leaf(name, child, key)
             s = s + sum(abs2, child)
+            n = n + length(child)
         else
-            s = s + _weight_l2(child, key)
+            cs, cn = _weight_l2_stats(child, key)
+            s = s + cs
+            n = n + cn
         end
     end
-    return s
+    return s, n
 end
diff --git a/test/test_extract_weights.jl b/test/test_extract_weights.jl
index 6029f75f..6c967a76 100644
--- a/test/test_extract_weights.jl
+++ b/test/test_extract_weights.jl
@@ -1,4 +1,5 @@
 using EasyHybrid
+using ComponentArrays
 using Lux
 using Random
 using Test
@@ -30,4 +31,12 @@ using Zygote
 
     # biases are not regularized by default
     @test weight_l2(ps; key = :bias) ≈ sum(sum(abs2, b) for b in extract_weights(ps; key = :bias))
+
+    n_weights = sum(length, ws)
+    @test weight_l2(ps; normalize = true) ≈ manual_l2 / n_weights
+
+    ps_ca = ComponentArray((Rb = ps,))
+    @test !isempty(extract_weights(ps_ca.Rb))
+    @test weight_l2(ps_ca.Rb) ≈ manual_l2
+    @test weight_l2(ps_ca.Rb; normalize = true) ≈ manual_l2 / n_weights
 end

From a407bd5027ed03148b949416b65ff26bae8dc153 Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Thu, 11 Jun 2026 16:35:47 +0200
Subject: [PATCH 13/16] Optimization.jl implementation with BFGS and LBFGS etc

---
 Project.toml                       |   2 +
 src/EasyHybrid.jl                  |   4 +-
 src/config/TrainingConfig.jl       |  24 +++-
 src/training/train_optimization.jl | 220 +++++++++++++++++++++++++++++
 src/training/training.jl           |   1 +
 5 files changed, 247 insertions(+), 4 deletions(-)
 create mode 100644 src/training/train_optimization.jl

diff --git a/Project.toml b/Project.toml
index e1888bc3..c469143d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -23,6 +23,7 @@ MLJ = "add582a8-e3ab-11e8-2d5e-e98b27df1bc7"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 NamedDims = "356022a1-0364-5f58-8944-0da4b18d706f"
+Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
@@ -63,6 +64,7 @@ MLUtils = "0.4.8"
 Makie = "0.22, 0.23, 0.24"
 NCDatasets = "0.14.8"
 NamedDims = "1.2.3"
+Optimization = "4, 5"
 OptimizationOptimisers = "0.3.7"
 OrderedCollections = "1.8.1"
 PrettyTables = "2.4.0, 3.1.2"
diff --git a/src/EasyHybrid.jl b/src/EasyHybrid.jl
index 99f43a3f..e29d53a4 100644
--- a/src/EasyHybrid.jl
+++ b/src/EasyHybrid.jl
@@ -32,6 +32,7 @@ using Lux: Lux
 using MLJ: partition
 using MLUtils: MLUtils, DataLoader, kfolds, numobs, rpad, splitobs
 using NCDatasets: NCDatasets, NCDataset, close, name
+using Optimization: Optimization, OptimizationFunction, OptimizationProblem, solve, remake
 using OptimizationOptimisers: OptimizationOptimisers, AdamW, Adam, Optimisers
 using OrderedCollections: OrderedDict
 using PrettyTables: PrettyTables
@@ -54,15 +55,16 @@ using Static: False, True
     using DataFrames
     using CSV
     using AxisKeys: axiskeys
+    using Optimization: Optimization, OptimizationFunction, OptimizationProblem, solve, remake
     using OptimizationOptimisers: OptimizationOptimisers, Optimisers, Adam, AdamW, RMSProp
     using ComponentArrays: ComponentArrays, ComponentArray
 end
 
 abstract type EasyHybridModels end
 
-include("config/config.jl")
 include("utils/utils.jl")
 include("models/models.jl")
+include("config/config.jl")
 include("data/data.jl")
 include("losses/losses.jl")
 include("training/training.jl")
diff --git a/src/config/TrainingConfig.jl b/src/config/TrainingConfig.jl
index 82a9360e..cb3a19cc 100644
--- a/src/config/TrainingConfig.jl
+++ b/src/config/TrainingConfig.jl
@@ -140,11 +140,23 @@ loss computation, data handling, output, and visualization.
     promote_f64::Bool = false
 
     """
-    `Optimization.jl` path only: build a validation `EpochSnapshot` (driving
-    history / early-stopping / dashboard) every `eval_every` callback hits.
-    Default: `1`.
+    `Optimization.jl` full-batch path only (`full_batch = true`): build a
+    validation `EpochSnapshot` (driving history / early-stopping / dashboard)
+    every `eval_every` solver iterations. Default: `1`.
     """
     eval_every::Int = 1
+
+    """
+    `Optimization.jl` minibatch path only (`full_batch = false`): number of
+    optimizer iterations to run on each *fixed* minibatch before resampling the
+    next one. Implements the "several L-BFGS steps per minibatch, then resample"
+    scheme of Le et al., 2011 (*On Optimization Methods for Deep Learning*, ICML,
+    §4.2): holding the minibatch fixed for a few iterations keeps the objective
+    (and L-BFGS curvature pairs / line search) consistent. `nepochs` controls the
+    number of outer resampling passes; `batchsize` the minibatch size. Setting
+    `inner_maxiters = 1` reduces to one optimizer step per minibatch. Default: `4`.
+    """
+    inner_maxiters::Int = 4
 end
 
 function validate_config(cfg::TrainConfig)
@@ -160,6 +172,12 @@ function validate_config(cfg::TrainConfig)
     cfg.patience > 0 ||
         throw(ArgumentError("patience must be positive, got $(cfg.patience)"))
 
+    cfg.eval_every > 0 ||
+        throw(ArgumentError("eval_every must be positive, got $(cfg.eval_every)"))
+
+    cfg.inner_maxiters > 0 ||
+        throw(ArgumentError("inner_maxiters must be positive, got $(cfg.inner_maxiters)"))
+
     check_training_loss(cfg.training_loss) # TODO: revisit implementation
 
     return cfg.training_loss in cfg.loss_types ||
diff --git a/src/training/train_optimization.jl b/src/training/train_optimization.jl
new file mode 100644
index 00000000..5dc9ffae
--- /dev/null
+++ b/src/training/train_optimization.jl
@@ -0,0 +1,220 @@
+"""
+    _train_optimization(model, data, train_cfg, data_cfg, solve_kwargs)
+
+`Optimization.jl`-based driver dispatched from `_train` whenever `train_cfg.opt`
+is not an `Optimisers.AbstractRule` (e.g. `Optim.LBFGS()` / `Optimization.LBFGS()`);
+see the
+[SciML minibatching tutorial](https://docs.sciml.ai/Optimization/stable/tutorials/minibatch/).
+
+Two modes, selected by `train_cfg.full_batch`:
+
+- `full_batch = true`: pass the full training set as a single tuple to one
+  `OptimizationProblem` and a single `solve(...)`. Batch-method idiom (the
+  recommended L-BFGS setup): the objective is a single consistent function.
+  `solve_kwargs` (e.g. `maxiters`, `g_abstol`, `f_reltol`) are splatted into
+  `solve`, and `train_cfg.eval_every` builds a validation `EpochSnapshot` every
+  N solver iterations via the callback.
+- `full_batch = false`: explicit "repeated minibatch" loop grounded in Le et
+  al., 2011 (*On Optimization Methods for Deep Learning*, ICML, §4.2). For each
+  of `train_cfg.nepochs` outer passes we iterate a (reshuffled) `DataLoader`
+  and, on each *fixed* minibatch, run `train_cfg.inner_maxiters` optimizer
+  iterations (`solve(...; maxiters = inner_maxiters)`), warm-starting the next
+  minibatch from the current parameters. Holding the minibatch fixed for a few
+  iterations keeps the objective (and L-BFGS curvature pairs / line search)
+  consistent — naive one-step-per-minibatch L-BFGS does not converge. A
+  validation `EpochSnapshot` is built once per outer pass. Optimization.jl's
+  own `DataLoader` iteration is **not** used here because it only applies to
+  the `Optimisers.jl`-style solvers, not Optim.jl's L-BFGS.
+
+Both modes honour `train_cfg.promote_f64`: promote `ps` to `Float64` before
+optimization (workaround for
+[Lux.jl#1260](https://github.com/LuxDL/Lux.jl/issues/1260)).
+"""
+function _train_optimization(model, data, train_cfg::TrainConfig, data_cfg::DataConfig, solve_kwargs::NamedTuple)
+    validate_config(train_cfg)
+    ext = load_makie_extension(train_cfg)
+    seed!(train_cfg.random_seed)
+
+    ((x_train, forcings_train), y_train), ((x_val, forcings_val), y_val) =
+        prepare_splits(data, model, data_cfg)
+    mask_train, _ = valid_mask(y_train)
+    mask_val, _ = valid_mask(y_val)
+
+    ps, st, _ = init_model_state(model, train_cfg)
+    ps_ca = ComponentArray(ps)
+    if train_cfg.promote_f64
+        ps_ca = ps_ca .|> Float64
+    end
+
+    init = compute_initial_state(
+        model, x_train, forcings_train, y_train, mask_train,
+        x_val, forcings_val, y_val, mask_val, ps_ca, st, train_cfg,
+    )
+    history = TrainingHistory(init)
+    stopper = EarlyStopping(init.l_val, ps_ca, st, train_cfg)
+    paths = resolve_paths(train_cfg)
+    prog = build_progress(train_cfg)
+    dashboard = init_dashboard(ext, init, train_cfg, y_train, y_val, model.targets)
+
+    save_initial_state!(paths, model, ps_ca, st, train_cfg)
+
+    loss_fn = _build_optim_loss(model, st, train_cfg)
+    opt_func = OptimizationFunction(loss_fn, train_cfg.autodiff_backend)
+
+    final_ps = Ref{Any}(ps_ca)
+    record_or_run(ext, paths, train_cfg) do io
+        if train_cfg.full_batch
+            # Convert once (not per objective/gradient evaluation): the full
+            # training set is the fixed `p` for the entire solve.
+            data_arg = collect_dim_data(
+                (x_train, forcings_train), (y_train, mask_train), train_cfg,
+            )
+            opt_prob = OptimizationProblem(opt_func, ps_ca, data_arg)
+            cb = _optim_callback(
+                model, st, init, history, stopper, dashboard, ext, prog, paths,
+                x_train, forcings_train, y_train, mask_train,
+                x_val, forcings_val, y_val, mask_val,
+                io, train_cfg, final_ps,
+            )
+            res = solve(opt_prob, train_cfg.opt; callback = cb, solve_kwargs...)
+            final_ps[] = res.u
+        else
+            _run_minibatch!(
+                final_ps, opt_func, ps_ca, model, st, init, history, stopper,
+                dashboard, ext, prog, paths,
+                x_train, forcings_train, y_train, mask_train,
+                x_val, forcings_val, y_val, mask_val,
+                io, train_cfg, solve_kwargs,
+            )
+        end
+    end
+
+    ps_out = final_ps[]
+    save_dashboard_img!(dashboard, ext, paths, train_cfg, stopper.best_epoch)
+    ps_out, st_out = best_or_final(stopper, ps_out, st, train_cfg)
+    save_final!(
+        paths, model, ps_out, st_out,
+        x_train, forcings_train, y_train,
+        x_val, forcings_val, y_val,
+        stopper, train_cfg,
+    )
+
+    return build_results(
+        model, history, stopper, ps_out, st_out,
+        x_train, forcings_train, y_train,
+        x_val, forcings_val, y_val,
+        train_cfg,
+    )
+end
+
+"""
+Build the scalar loss closure consumed by `Optimization.jl`, called as
+`loss_fn(p, data)`. `data` is an **already device-placed / `Array`-converted**
+batch (shape `((x, forcings), (y, mask))`), produced once per (mini)batch by
+`collect_dim_data` in the caller — *not* inside this closure. Keeping the data
+prep out of the loss is important: L-BFGS line searches call the objective many
+times per iteration, so re-running `collect_dim_data` (NamedTuple rebuilds,
+`Array` copies, `gdev` transfers) on every evaluation was a major slowdown,
+especially on the minibatch path. It also keeps the closure trivially
+Zygote-differentiable (no `pairs(...)`/`∇map` in the AD tape).
+"""
+function _build_optim_loss(model, st, cfg::TrainConfig)
+    logging = LoggingLoss(
+        train_mode = true,
+        loss_types = cfg.loss_types,
+        training_loss = cfg.training_loss,
+        extra_loss = cfg.extra_loss,
+        agg = cfg.agg,
+    )
+    return function (p, data)
+        loss_val, _, _ = compute_loss(model, p, st, data; logging)
+        return loss_val
+    end
+end
+
+"""
+Explicit "repeated minibatch" driver for the `full_batch = false` Optimization
+path (Le et al., 2011, ICML, §4.2). For each of `cfg.nepochs` outer passes we
+iterate a reshuffled `DataLoader`; on every *fixed* minibatch we run
+`cfg.inner_maxiters` optimizer iterations warm-started from the current `ps`,
+then resample. A validation `EpochSnapshot` (history / early-stopping /
+dashboard / checkpoint) is built once per outer pass, so `patience` is counted
+in outer passes — consistent with the `Optimisers.jl` loop.
+
+`maxiters` / `epochs` from `solve_kwargs` are dropped (the per-minibatch budget
+is `cfg.inner_maxiters` and the pass count is `cfg.nepochs`); any remaining
+`solve_kwargs` (e.g. `g_abstol`, `f_reltol`) are forwarded to each inner solve.
+"""
+function _run_minibatch!(
+        final_ps, opt_func, ps0, model, st, init, history, stopper,
+        dashboard, ext, prog, paths,
+        x_train, forcings_train, y_train, mask_train,
+        x_val, forcings_val, y_val, mask_val,
+        io, cfg::TrainConfig, solve_kwargs::NamedTuple,
+    )
+    loader = build_loader(x_train, forcings_train, y_train, mask_train, cfg)
+    inner_kwargs = Base.structdiff(solve_kwargs, (; maxiters = nothing, epochs = nothing))
+    ps = ps0
+
+    # Build the problem once and reuse it across minibatches via `remake`, which
+    # recreates it cheaply with a new initial guess (`u0`) and minibatch data
+    # (`p`) while keeping the same `OptimizationFunction` — see the SciML
+    # polyalgorithm tutorial. The placeholder batch is immediately replaced.
+    x0b, y0b = first(loader)
+    opt_prob = OptimizationProblem(opt_func, ps, collect_dim_data(x0b, y0b, cfg))
+
+    for epoch in 1:cfg.nepochs
+        for (x, y) in loader
+            isemptybatch(y[2]) && continue
+            # Convert the minibatch once here, then hold it fixed as `p` for all
+            # `inner_maxiters` solver iterations — avoids re-running
+            # `collect_dim_data` on every line-search objective evaluation.
+            data = collect_dim_data(x, y, cfg)
+            opt_prob = remake(opt_prob; u0 = ps, p = data)
+            res = solve(opt_prob, cfg.opt; maxiters = cfg.inner_maxiters, inner_kwargs...)
+            ps = res.u
+            final_ps[] = ps
+        end
+
+        snapshot = evaluate_epoch(
+            model, x_train, forcings_train, y_train, mask_train,
+            x_val, forcings_val, y_val, mask_val,
+            ps, st, init, cfg,
+        )
+        update!(stopper, history, snapshot, ps, st, epoch, cfg)
+        save_epoch!(paths, model, ps, st, snapshot, epoch, cfg)
+        update_dashboard!(dashboard, ext, snapshot, epoch, io, cfg)
+        log_progress!(prog, init, snapshot, epoch, cfg)
+
+        is_done(stopper) && break
+    end
+
+    return final_ps
+end
+
+function _optim_callback(
+        model, st, init, history, stopper, dashboard, ext, prog, paths,
+        x_train, forcings_train, y_train, mask_train,
+        x_val, forcings_val, y_val, mask_val,
+        io, cfg::TrainConfig, final_ps,
+    )
+    return function (state, _loss)
+        iter = state.iter
+        ps_cur = state.u
+        final_ps[] = ps_cur
+
+        if iter > 0 && iter % cfg.eval_every == 0
+            snapshot = evaluate_epoch(
+                model, x_train, forcings_train, y_train, mask_train,
+                x_val, forcings_val, y_val, mask_val,
+                ps_cur, st, init, cfg,
+            )
+            update!(stopper, history, snapshot, ps_cur, st, iter, cfg)
+            save_epoch!(paths, model, ps_cur, st, snapshot, iter, cfg)
+            update_dashboard!(dashboard, ext, snapshot, iter, io, cfg)
+            log_progress!(prog, init, snapshot, iter, cfg)
+        end
+
+        return is_done(stopper)
+    end
+end
diff --git a/src/training/training.jl b/src/training/training.jl
index 558aeb57..8e0eb868 100644
--- a/src/training/training.jl
+++ b/src/training/training.jl
@@ -4,5 +4,6 @@ include("epoch.jl")
 include("dashboard.jl")
 include("early_stopping.jl")
 include("train.jl")
+include("train_optimization.jl")
 include("show_train.jl")
 include("tune.jl")

From 50fcc924b745dca1282fafb45365efa435f3e61f Mon Sep 17 00:00:00 2001
From: Bernhard Ahrens <bahrens@bgc-jena.mpg.de>
Date: Thu, 11 Jun 2026 16:36:27 +0200
Subject: [PATCH 14/16] warning of train test for evaluation

---
 src/training/train.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/training/train.jl b/src/training/train.jl
index 62840bfd..25c5dab4 100644
--- a/src/training/train.jl
+++ b/src/training/train.jl
@@ -345,6 +345,11 @@ function rename_deprecated_kwargs(kwargs)
 end
 
 function evaluate_acc(ghm, x, forcings, y, y_no_nan, ps, st, loss_types, training_loss, extra_loss, agg)
+    # Metric/validation evaluation is a plain forward pass (no autodiff). Switch
+    # to test mode so layers like BatchNorm use their running statistics instead
+    # of batch statistics, and to avoid LuxLib's "training is set to Val{true}()
+    # but is not being used within an autodiff call" warning.
+    st = LuxCore.testmode(st)
     loss_val, sts, ŷ = compute_loss(ghm, ps, st, ((x, forcings), (y, y_no_nan)), logging = LoggingLoss(train_mode = false, loss_types = loss_types, training_loss = training_loss, extra_loss = extra_loss, agg = agg))
     return loss_val, sts, ŷ
 end

From aafa4907916e2a0a58b7f13c65c59da1e8fbc15c Mon Sep 17 00:00:00 2001
From: BernhardAhrens <32769547+BernhardAhrens@users.noreply.github.com>
Date: Fri, 12 Jun 2026 12:52:53 +0200
Subject: [PATCH 15/16] runic

---
 src/training/train.jl        | 18 +++++++++---------
 src/utils/extract_weights.jl |  4 ++--
 src/utils/utils.jl           |  2 +-
 3 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/src/training/train.jl b/src/training/train.jl
index 25c5dab4..a7f53989 100644
--- a/src/training/train.jl
+++ b/src/training/train.jl
@@ -209,11 +209,11 @@ A [`TrainResults`](@ref) with the following fields:
 - `best_loss`: Best validation loss recorded during training.
 """
 function train(
-    model, data;
-    train_cfg::TrainConfig = TrainConfig(),
-    data_cfg::DataConfig  = DataConfig(),
-    kwargs...,
-)
+        model, data;
+        train_cfg::TrainConfig = TrainConfig(),
+        data_cfg::DataConfig = DataConfig(),
+        kwargs...,
+    )
     train_cfg, data_cfg, solve_kwargs = override_configs(train_cfg, data_cfg, kwargs)
     return _train(model, data, train_cfg, data_cfg, solve_kwargs)
 end
@@ -305,12 +305,12 @@ function override_configs(train_cfg::TrainConfig, data_cfg::DataConfig, kwargs)
     kwargs = expand_sequence_kwargs(kwargs)
 
     train_overrides = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k in train_keys)
-    data_overrides  = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k in data_keys)
-    solve_kwargs    = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k ∉ train_keys && k ∉ data_keys)
+    data_overrides = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k in data_keys)
+    solve_kwargs = NamedTuple(k => kwargs[k] for k in keys(kwargs) if k ∉ train_keys && k ∉ data_keys)
 
     return override_config(train_cfg, train_overrides),
-           override_config(data_cfg, data_overrides),
-           solve_kwargs
+        override_config(data_cfg, data_overrides),
+        solve_kwargs
 end
 
 """
diff --git a/src/utils/extract_weights.jl b/src/utils/extract_weights.jl
index 78de9727..6c8f401b 100644
--- a/src/utils/extract_weights.jl
+++ b/src/utils/extract_weights.jl
@@ -71,10 +71,10 @@ function weight_l2(ps; key::Symbol = :weight, normalize::Bool = false)
     return normalize ? (n > 0 ? s / n : zero(s)) : s
 end
 
-_weight_l2_stats(::Any, ::Symbol) = (0f0, 0)
+_weight_l2_stats(::Any, ::Symbol) = (0.0f0, 0)
 
 function _weight_l2_stats(node::Union{NamedTuple, ComponentArray}, key::Symbol)
-    s = 0f0
+    s = 0.0f0
     n = 0
     for name in propertynames(node)
         child = getproperty(node, name)
diff --git a/src/utils/utils.jl b/src/utils/utils.jl
index 1389bb97..5d84e255 100644
--- a/src/utils/utils.jl
+++ b/src/utils/utils.jl
@@ -5,4 +5,4 @@ include("plotrecipes.jl")
 include("helpers_data_loading.jl")
 include("helpers_cross_validation.jl")
 include("print_banner.jl")
-include("extract_weights.jl")
\ No newline at end of file
+include("extract_weights.jl")

From 0c1b1c6169b0f05e9ca096d39e55c8aaabb31205 Mon Sep 17 00:00:00 2001
From: BernhardAhrens <32769547+BernhardAhrens@users.noreply.github.com>
Date: Fri, 12 Jun 2026 12:54:47 +0200
Subject: [PATCH 16/16] fix tests

---
 test/Project.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/Project.toml b/test/Project.toml
index 9e472887..74824f81 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -13,6 +13,7 @@ Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [sources]