Fix gradient-aggregation bias in GradientMode parameter updates

Project.toml

@@ -1,7 +1,7 @@
 name = "ApplicationDrivenLearning"
 uuid = "0856f1c8-ef17-4e14-9230-2773e47a789e"
 authors = ["Giovanni Amorim", "Joaquim Garcia"]
-version = "0.1.6"
+version = "0.1.7"
 
 [deps]
 BilevelJuMP = "485130c0-026e-11ea-0f1a-6992cd14145c"
@@ -27,6 +27,7 @@ JuMP = "1.24"
 MPI = "0.20.22"
 Optim = "1.11"
 Optimisers = "0.4.5"
-ParametricOptInterface = "0.9.0"
-Zygote = "0.6.75"
+ParametricOptInterface = "0.9.0, 0.10"
+Statistics = "1"
+Zygote = "0.6.75, 0.7"
 julia = "1.10"

docs/src/tutorials/modes.md

@@ -223,7 +223,7 @@ rules from Flux.jl.
 after a specified number of epochs. This enables faster iterations with the drawback of
 possibly missing sets of parameters with low associated cost.
 - `time_limit`: The time limit for the gradient algorithm in seconds.
-- `g_tol`: Convergence condition on the infinite norm of the gradient vector. Below, we illustrate the use of NelderMeadMode to optimize the predictive model used in the ongoing example.
+- `g_tol`: Convergence condition on the infinity norm of the per-sample cost gradients with respect to the forecasts (the maximum absolute entry over all samples and outputs). Below, we illustrate the use of NelderMeadMode to optimize the predictive model used in the ongoing example.
 
 ### Example
 

src/optimizers/gradient_mpi.jl

@@ -75,8 +75,8 @@ function train_with_gradient_mpi!(
                     (v) -> compute_cost_and_gradients(v[1], v[2], true),
                     [[curr_θ, i] for i in batches[epoch, :]],
                 )
-                dCdy =
-                    sum([r[2] for r in pmap_result_with_gradients]) ./ batch_size
+                # stack per-sample gradients into rows aligned with `epochx`
+                dC = reduce(vcat, [r[2]' for r in pmap_result_with_gradients])
 
                 if compute_full_cost
                     # broadcast `is_done = false` again
@@ -100,7 +100,8 @@ function train_with_gradient_mpi!(
                     [[curr_θ, i] for i = 1:T],
                 )
                 curr_C = sum([r[1] for r in pmap_result]) ./ T
-                dCdy = sum([r[2] for r in pmap_result]) ./ T
+                # stack per-sample gradients into rows aligned with `epochx`
+                dC = reduce(vcat, [r[2]' for r in pmap_result])
             end
 
             if compute_full_cost
@@ -129,15 +130,15 @@ function train_with_gradient_mpi!(
             end
 
             # check gradient tolerance
-            if maximum(abs.(dCdy)) < g_tol
+            if maximum(abs.(dC)) < g_tol
                 if verbose
                     println("Gradient tolerance reached.")
                 end
                 break
             end
 
             # take gradient step (if not last epoch)
-            apply_gradient!(model.forecast, dCdy, epochx, opt_state)
+            apply_gradient!(model.forecast, dC, epochx, opt_state)
         end
 
         # release workers

src/predictive_model.jl

@@ -331,27 +331,34 @@ function apply_params(model::PredictiveModel, θ)
 end
 
 """
-    apply_gradient!(model, dCdy, X, rule)
+    apply_gradient!(model, dCdy, X, opt_state)
 
-Apply a gradient vector to the model parameters.
+Apply per-sample cost gradients to the model parameters.
+
+The optimization layer provides `dCdy`, the gradient of the assessment cost with
+respect to the forecasts, for each sample. Because the optimization step itself
+is not differentiable by the AD backend, these gradients are propagated through
+the forecast model with a linear surrogate loss whose parameter-gradient equals
+the chain-rule term `(1/T) Σₜ (dC/dŷₜ)·(dŷₜ/dθ)`.
 
 ...
 
 # Arguments
 
   - `model::PredictiveModel`: model to be updated.
-  - `dCdy::Vector{<:Real}`: gradient vector.
-  - `X::Matrix{<:Real}`: input data.
-  - `rule`: Optimisation rule.
+  - `dCdy::AbstractMatrix{<:Real}`: per-sample cost gradients, size
+    `(T, output_size)`, with row `t` aligned to sample `t` (row `t` of `X`).
+  - `X::Matrix{<:Real}`: input data, size `(T, input_size)`.
+  - `opt_state`: Optimisers optimisation state.
     ...
 """
 function apply_gradient!(
     model::PredictiveModel,
-    dCdy::AbstractVector{<:Real},
+    dCdy::AbstractMatrix{<:Real},
     X::Matrix{<:Real},
     opt_state,
 )
-    loss3(m, X) = mean(dCdy'm(X'))
+    loss3(m, X) = sum(dCdy' .* m(X')) / size(X, 1)
     grad = Zygote.gradient(loss3, model, X)[1]
     return Optimisers.update!(opt_state, model, grad)
 end

src/simulation.jl

@@ -75,7 +75,11 @@ Compute the cost function (C) based on the model predictions and the true values
   - `model::ApplicationDrivenLearning.Model`: model to evaluate.
   - `X::Matrix{<:Real}`: input data.
   - `Y::Matrix{<:Real}`: true values.
-  - `with_gradients::Bool=false`: flag to compute and return gradients.
+  - `with_gradients::Bool=false`: flag to compute and return the cost gradients
+    with respect to the forecasts. When set, the second returned value is the
+    per-sample gradient matrix of size `(T, output_size)`.
+  - `aggregate::Bool=true`: when true, the returned cost is averaged over the `T`
+    samples. Only affects the cost; gradients are always per-sample.
     ...
 """
 function compute_cost(
@@ -122,7 +126,6 @@ function compute_cost(
     # aggregate cost if requested
     if aggregate
         C = sum(C) / T
-        dC = sum(dC, dims = 1)[1, :] / T
     end
 
     if with_gradients

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+ApplicationDrivenLearning = "0856f1c8-ef17-4e14-9230-2773e47a789e"
 BilevelJuMP = "485130c0-026e-11ea-0f1a-6992cd14145c"
 DiffOpt = "930fe3bc-9c6b-11ea-2d94-6184641e85e7"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"

test/test_gradient.jl

@@ -39,3 +39,92 @@ Y = Dict(d => Float32.(ones(1)))
     sol = ApplicationDrivenLearning.train!(model, X, Y, opt)
     @test initial_sol == sol.params
 end
+
+@testset "GradientMode Multi-Sample Gradient Correctness" begin
+    # Newsvendor model where the per-sample cost-gradient (dC/dŷ) depends on
+    # whether the forecast over- or under-shoots the realized demand. The
+    # gradient that GradientMode applies to the parameters must equal the true
+    # gradient of the aggregate cost, dC/dθ = (1/T) Σ_t (dC/dŷ_t)·(dŷ_t/dθ).
+    #
+    # The samples below are chosen so the over/under-stock regime correlates
+    # with the input, i.e. mean_t(dCdy_t · J_t) ≠ mean_t(dCdy_t)·mean_t(J_t).
+    # Aggregating dCdy across the batch *before* forming the surrogate loss
+    # (mean(dCdy' * m(X'))) computes the right-hand side instead of the
+    # left-hand side and is therefore biased for T > 1. This test pins the
+    # left-hand side via finite differences of the true cost.
+    c, q, r = 1.0, 3.0, 0.0
+    nv = ApplicationDrivenLearning.Model()
+    @variables(nv, begin
+        x, ApplicationDrivenLearning.Policy
+        d, ApplicationDrivenLearning.Forecast
+    end)
+    @variables(ApplicationDrivenLearning.Plan(nv), begin
+        yp >= 0
+        wp >= 0
+    end)
+    @constraints(ApplicationDrivenLearning.Plan(nv), begin
+        yp <= d.plan
+        yp + wp <= x.plan
+    end)
+    @objective(
+        ApplicationDrivenLearning.Plan(nv),
+        Min,
+        c * x.plan - q * yp - r * wp
+    )
+    @variables(ApplicationDrivenLearning.Assess(nv), begin
+        ya >= 0
+        wa >= 0
+    end)
+    @constraints(ApplicationDrivenLearning.Assess(nv), begin
+        ya <= d.assess
+        ya + wa <= x.assess
+    end)
+    @objective(
+        ApplicationDrivenLearning.Assess(nv),
+        Min,
+        c * x.assess - q * ya - r * wa
+    )
+    set_optimizer(nv, HiGHS.Optimizer)
+    set_silent(nv)
+    ApplicationDrivenLearning.set_forecast_model(
+        nv,
+        Chain(Dense(1 => 1; bias = false, init = (s...) -> ones(s...))),
+    )
+
+    # x = [1,2,3], θ₀ = 1 ⇒ ŷ = [1,2,3]; d = [5,5,1] ⇒ samples 1,2 understock
+    # (dCdy = c-q = -2), sample 3 overstocks (dCdy = c = +1).
+    # True gradient = mean(dCdy_t · x_t) = (-2·1 -2·2 +1·3)/3 = -1.
+    # Aggregated/biased gradient = mean(dCdy_t)·mean(x_t) = (-1)·2 = -2.
+    Xnv = reshape([1.0, 2.0, 3.0], 3, 1)
+    Ynv = reshape([5.0, 5.0, 1.0], 3, 1)
+
+    θ0 = ApplicationDrivenLearning.extract_params(nv.forecast)
+
+    # Reference: central finite-difference gradient of the true aggregate cost.
+    function cost_at(θ)
+        ApplicationDrivenLearning.apply_params(nv.forecast, θ)
+        return ApplicationDrivenLearning.compute_cost(nv, Xnv, Ynv)
+    end
+    h = 1e-4
+    fd_grad = similar(θ0)
+    for i in eachindex(θ0)
+        θp = copy(θ0)
+        θm = copy(θ0)
+        θp[i] += h
+        θm[i] -= h
+        fd_grad[i] = (cost_at(θp) - cost_at(θm)) / (2h)
+    end
+    ApplicationDrivenLearning.apply_params(nv.forecast, θ0)
+
+    # Gradient actually applied by the framework, recovered from a single
+    # Descent(η) step: θ₁ = θ₀ - η·g  ⇒  g = (θ₀ - θ₁)/η.
+    η = 1.0
+    opt_state = Flux.setup(Flux.Descent(η), nv.forecast)
+    _, dC = ApplicationDrivenLearning.compute_cost(nv, Xnv, Ynv, true)
+    ApplicationDrivenLearning.apply_gradient!(nv.forecast, dC, Xnv, opt_state)
+    θ1 = ApplicationDrivenLearning.extract_params(nv.forecast)
+    applied_grad = (θ0 .- θ1) ./ η
+    ApplicationDrivenLearning.apply_params(nv.forecast, θ0)
+
+    @test applied_grad ≈ fd_grad atol = 1e-2
+end

test/test_predictive_model.jl

@@ -22,7 +22,7 @@ out_size = 2
 
     ApplicationDrivenLearning.apply_gradient!(
         forecaster,
-        ones(out_size),
+        ones((1, out_size)),
         ones((1, in_size)),
         Flux.setup(Flux.Descent(0.1), forecaster),
     )
@@ -50,7 +50,7 @@ end
 
     ApplicationDrivenLearning.apply_gradient!(
         forecaster,
-        ones(out_size),
+        ones((1, out_size)),
         ones((1, in_size)),
         Flux.setup(Flux.Descent(0.1), forecaster),
     )
@@ -86,7 +86,7 @@ end
 
     ApplicationDrivenLearning.apply_gradient!(
         forecaster,
-        ones(out_size),
+        ones((1, out_size)),
         ones((1, in_size)),
         Flux.setup(Flux.Descent(0.1), forecaster),
     )
@@ -126,7 +126,7 @@ end
 
     ApplicationDrivenLearning.apply_gradient!(
         forecaster,
-        ones(out_size),
+        ones((1, out_size)),
         ones((1, in_size)),
         Flux.setup(Flux.Descent(0.1), forecaster),
     )