Make Frank-Wolfe regularized predictor compatible with generic arrays

Author: gdalle

src/frank_wolfe/differentiable_frank_wolfe.jl

@@ -27,12 +27,14 @@ end
 ## Forward pass
 
 function optimal_active_set(
-    dfw::DifferentiableFrankWolfe, θ::AbstractArray{<:Real}; fw_kwargs=(;)
+    dfw::DifferentiableFrankWolfe,
+    θ::AbstractArray{<:Real},
+    x0::AbstractArray{<:Real};
+    fw_kwargs=(;),
 )
     (; f, ∇ₓf, lmo) = dfw
     obj(x) = f(x, θ)
     grad!(g, x) = g .= ∇ₓf(x, θ)
-    x0 = compute_extreme_point(lmo, zero(θ))
     full_fw_kwargs = merge(DEFAULT_FRANK_WOLFE_KWARGS, fw_kwargs)
     x, v, primal, dual_gap, traj_data, active_set = away_frank_wolfe(
         obj, grad!, lmo, x0; full_fw_kwargs...
@@ -41,62 +43,66 @@ function optimal_active_set(
     return active_set
 end
 
-function (dfw::DifferentiableFrankWolfe)(θ::AbstractArray{<:Real}; fw_kwargs=(;))
-    active_set::ActiveSet = optimal_active_set(dfw, θ; fw_kwargs=fw_kwargs)
+function (dfw::DifferentiableFrankWolfe)(
+    θ::AbstractArray{<:Real}, x0::AbstractArray{<:Real}; fw_kwargs=(;)
+)
+    active_set::ActiveSet = optimal_active_set(dfw, θ, x0; fw_kwargs=fw_kwargs)
     return active_set.x
 end
 
-## Backward pass, only works with vectors
+## Backward pass
 
 function frank_wolfe_optimality_conditions(
     dfw::DifferentiableFrankWolfe,
-    θ::AbstractVector{<:Real},
+    θ::AbstractArray{<:Real},
     p::AbstractVector{<:Real},
-    V::AbstractMatrix{<:Real},
+    A::AbstractVector{<:AbstractArray{<:Real}},
 )
     (; ∇ₓf) = dfw
-    ∇ₚg = V' * ∇ₓf(V * p, θ)
+    x = sum(pᵢ * Aᵢ for (pᵢ, Aᵢ) in zip(p, A))
+    b = ∇ₓf(x, θ)
+    ∇ₚg = [dot(Aᵢ, b) for Aᵢ in A]
     T = sparse_argmax(p - ∇ₚg)
     return T - p
 end
 
 function ChainRulesCore.rrule(
-    rc::RuleConfig, dfw::DifferentiableFrankWolfe, θ::AbstractVector{<:Real}; fw_kwargs=(;)
-)
+    rc::RuleConfig,
+    dfw::DifferentiableFrankWolfe,
+    θ::AbstractArray{R1},
+    x0::AbstractArray{R2};
+    fw_kwargs=(;),
+) where {R1<:Real,R2<:Real}
+    R = promote_type(R1, R2)
     (; linear_solver) = dfw
 
-    active_set::ActiveSet = optimal_active_set(dfw, θ; fw_kwargs=fw_kwargs)
-    V = reduce(hcat, active_set.atoms)
+    active_set::ActiveSet = optimal_active_set(dfw, θ, x0; fw_kwargs=fw_kwargs)
+    A = active_set.atoms
     p = active_set.weights
-    n, m = length(θ), length(p)
+    x = active_set.x
 
-    conditions_θ(θ_bis) = frank_wolfe_optimality_conditions(dfw, θ_bis, p, V)
-    conditions_p(p_bis) = -frank_wolfe_optimality_conditions(dfw, θ, p_bis, V)
+    conditions_θ(θ_bis) = frank_wolfe_optimality_conditions(dfw, θ_bis, p, A)
+    conditions_p(p_bis) = -frank_wolfe_optimality_conditions(dfw, θ, p_bis, A)
 
     pullback_Aᵀ = last ∘ rrule_via_ad(rc, conditions_p, p)[2]
     pullback_Bᵀ = last ∘ rrule_via_ad(rc, conditions_θ, θ)[2]
 
-    mul_Aᵀ!(res, v) = res .= pullback_Aᵀ(v)
-    mul_Bᵀ!(res, v) = res .= pullback_Bᵀ(v)
+    mul_Aᵀ!(res, u::AbstractVector) = res .= vec(pullback_Aᵀ(reshape(u, size(p))))
+    mul_Bᵀ!(res, v::AbstractVector) = res .= vec(pullback_Bᵀ(reshape(v, size(p))))
 
-    Aᵀ = LinearOperator(Float64, m, m, false, false, mul_Aᵀ!)
-    Bᵀ = LinearOperator(Float64, n, m, false, false, mul_Bᵀ!)
+    n, m = length(θ), length(p)
+    Aᵀ = LinearOperator(R, m, m, false, false, mul_Aᵀ!)
+    Bᵀ = LinearOperator(R, n, m, false, false, mul_Bᵀ!)
 
     function frank_wolfe_pullback(dx)
-        dp = V' * Vector(unthunk(dx))
+        dx = unthunk(dx)
+        dp = [dot(Aᵢ, dx) for Aᵢ in A]
         u, stats = linear_solver(Aᵀ, dp)
-        if !stats.solved
-            error("The linear solver failed to converge")
-        end
-        dθ = Bᵀ * u
-        return (NoTangent(), dθ)
+        stats.solved || error("Linear solver failed to converge")
+        dθ_vec = Bᵀ * u
+        dθ = reshape(dθ_vec, size(θ))
+        return (NoTangent(), dθ, NoTangent())
     end
 
-    return active_set.x, frank_wolfe_pullback
-end
-
-function ChainRulesCore.rrule(
-    rc::RuleConfig, dfw::DifferentiableFrankWolfe, θ::AbstractArray{<:Real}; fw_kwargs=(;)
-)
-    throw(ArgumentError("θ must be a vector and not a higher-dimensional array"))
+    return x, frank_wolfe_pullback
 end

src/regularized/regularized_generic.jl

@@ -1,38 +1,38 @@
 """
-    RegularizedGeneric{RF,RG,F,G,M,S}
+    RegularizedGeneric{M,RF,RG,F,G,S}
 
 Generic and differentiable regularized prediction function `ŷ(θ) = argmax {θᵀy - Ω(y)}`.
 
 Relies on the Frank-Wolfe algorithm to minimize a concave objective on a polytope.
 
 # Fields
+- `maximizer::M`
 - `Ω::RF`
 - `∇Ω::RG`
 - `f::F`
 - `∇ₓf::G`
-- `maximizer::M`
 - `linear_solver::S`
 
 See also: [`DifferentiableFrankWolfe`](@ref).
 """
-struct RegularizedGeneric{RF,RG,F,G,M,S}
+struct RegularizedGeneric{M,RF,RG,F,G,S}
+    maximizer::M
     Ω::RF
     ∇Ω::RG
     f::F
     ∇ₓf::G
-    maximizer::M
     linear_solver::S
 end
 
 function Base.show(io::IO, regularized::RegularizedGeneric)
-    (; Ω, ∇Ω, maximizer, linear_solver) = regularized
-    return print(io, "RegularizedGeneric($Ω, $∇Ω, $maximizer, $linear_solver)")
+    (; maximizer, Ω, ∇Ω, linear_solver) = regularized
+    return print(io, "RegularizedGeneric($maximizer, $Ω, $∇Ω, $linear_solver)")
 end
 
-function RegularizedGeneric(Ω, ∇Ω, maximizer; linear_solver=gmres)
+function RegularizedGeneric(maximizer, Ω, ∇Ω; linear_solver=gmres)
     f(y, θ) = Ω(y) - dot(θ, y)
     ∇ₓf(y, θ) = ∇Ω(y) - θ
-    return RegularizedGeneric(Ω, ∇Ω, f, ∇ₓf, maximizer, linear_solver)
+    return RegularizedGeneric(maximizer, Ω, ∇Ω, f, ∇ₓf, linear_solver)
 end
 
 @traitimpl IsRegularized{RegularizedGeneric}
@@ -43,15 +43,6 @@ end
 
 ## Forward pass
 
-function (regularized::RegularizedGeneric)(
-    θ::AbstractArray{<:Real}; maximizer_kwargs=(;), fw_kwargs=(;)
-)
-    (; f, ∇ₓf, maximizer, linear_solver) = regularized
-    lmo = LMOWrapper(maximizer, maximizer_kwargs)
-    dfw = DifferentiableFrankWolfe(f, ∇ₓf, lmo, linear_solver)
-    return dfw(θ; fw_kwargs=fw_kwargs)
-end
-
 function optimal_active_set(
     regularized::RegularizedGeneric,
     θ::AbstractArray{<:Real};
@@ -61,30 +52,33 @@ function optimal_active_set(
     (; f, ∇ₓf, maximizer, linear_solver) = regularized
     lmo = LMOWrapper(maximizer, maximizer_kwargs)
     dfw = DifferentiableFrankWolfe(f, ∇ₓf, lmo, linear_solver)
-    return optimal_active_set(dfw, θ; fw_kwargs=fw_kwargs)
+    x0 = compute_extreme_point(lmo, θ)
+    return optimal_active_set(dfw, θ, x0; fw_kwargs=fw_kwargs)
+end
+
+function (regularized::RegularizedGeneric)(
+    θ::AbstractArray{<:Real}; maximizer_kwargs=(;), fw_kwargs=(;)
+)
+    active_set = optimal_active_set(
+        regularized, θ; maximizer_kwargs=maximizer_kwargs, fw_kwargs=fw_kwargs
+    )
+    return active_set.x
 end
 
 ## Backward pass, only works with vectors
 
 function ChainRulesCore.rrule(
     rc::RuleConfig,
     regularized::RegularizedGeneric,
-    θ::AbstractVector{<:Real};
+    θ::AbstractArray{<:Real};
     maximizer_kwargs=(;),
     fw_kwargs=(;),
 )
     (; f, ∇ₓf, maximizer, linear_solver) = regularized
     lmo = LMOWrapper(maximizer, maximizer_kwargs)
     dfw = DifferentiableFrankWolfe(f, ∇ₓf, lmo, linear_solver)
-    return rrule(rc, dfw, θ; fw_kwargs=fw_kwargs)
-end
-
-function ChainRulesCore.rrule(
-    rc::RuleConfig,
-    regularized::RegularizedGeneric,
-    θ::AbstractArray{<:Real};
-    maximizer_kwargs=(;),
-    fw_kwargs=(;),
-)
-    throw(ArgumentError("θ must be a vector and not a higher-dimensional array"))
+    x0 = compute_extreme_point(lmo, θ)
+    x, frank_wolfe_pullback = rrule(rc, dfw, θ, x0; fw_kwargs=fw_kwargs)
+    regularized_generic_pullback(dx) = frank_wolfe_pullback(dx)[1:2]
+    return x, regularized_generic_pullback
 end

test/argmax.jl

@@ -62,13 +62,13 @@ pipelines_imitation_y = [
         encoder=encoder_factory(),
         maximizer=identity,
         loss=FenchelYoungLoss(
-            RegularizedGeneric(half_square_norm, identity, true_maximizer)
+            RegularizedGeneric(true_maximizer, half_square_norm, identity)
         ),
     ),
     # Generic regularized + other loss
     (
         encoder=encoder_factory(),
-        maximizer=RegularizedGeneric(half_square_norm, identity, true_maximizer),
+        maximizer=RegularizedGeneric(true_maximizer, half_square_norm, identity),
         loss=Flux.Losses.mse,
     ),
 ]

test/frank_wolfe.jl

@@ -8,6 +8,7 @@ using Zygote
 Random.seed!(63)
 
 d = 100
+x0 = ones(d) / d;
 θ = rand(d);
 v = rand(d);
 rc = Zygote.ZygoteRuleConfig()
@@ -23,20 +24,20 @@ f(x, θ) = half_square_norm(x - θ)
 lmo = FrankWolfe.UnitSimplexOracle(1.0)
 
 dfw = DifferentiableFrankWolfe(f, ∇ₓf, lmo)
-_, pullback_dfw = rrule_via_ad(rc, dfw, θ; fw_kwargs=fw_kwargs);
+_, pullback_dfw = rrule_via_ad(rc, dfw, θ, x0; fw_kwargs=fw_kwargs);
 
 @testset verbose = true "DifferentiableFrankWolfe" begin
-    @test mean(abs, dfw(θ; fw_kwargs=fw_kwargs) - sparse_argmax(θ)) < 1e-3
+    @test mean(abs, dfw(θ, x0; fw_kwargs=fw_kwargs) - sparse_argmax(θ)) < 1e-3
     @test mean(abs, pullback_dfw(v)[2] - pullback_sparse_argmax(v)[2]) < 1e-3
 end
 
 ## RegularizedGeneric
 
+maximizer(θ) = one_hot_argmax(θ)
 Ω(y) = half_square_norm(y)
 ∇Ω(y) = y
-maximizer(θ) = one_hot_argmax(θ)
 
-regularized = RegularizedGeneric(Ω, ∇Ω, maximizer)
+regularized = RegularizedGeneric(maximizer, Ω, ∇Ω)
 _, pullback_regularized = rrule_via_ad(rc, regularized, θ; fw_kwargs=fw_kwargs);
 
 @testset verbose = true "RegularizedGeneric" begin

test/paths.jl

@@ -58,7 +58,7 @@ pipelines_imitation_y = [
         encoder=encoder_factory(),
         maximizer=identity,
         loss=FenchelYoungLoss(
-            RegularizedGeneric(half_square_norm, identity, true_maximizer)
+            RegularizedGeneric(true_maximizer, half_square_norm, identity)
         ),
     ),
 ]

test/ranking.jl

@@ -56,13 +56,13 @@ pipelines_imitation_y = [
         encoder=encoder_factory(),
         maximizer=identity,
         loss=FenchelYoungLoss(
-            RegularizedGeneric(half_square_norm, identity, true_maximizer)
+            RegularizedGeneric(true_maximizer, half_square_norm, identity)
         ),
     ),
     # Generic regularized + other loss
     (
         encoder=encoder_factory(),
-        maximizer=RegularizedGeneric(half_square_norm, identity, true_maximizer),
+        maximizer=RegularizedGeneric(true_maximizer, half_square_norm, identity),
         loss=Flux.Losses.mse,
     ),
 ]

test/runtests.jl

@@ -18,9 +18,6 @@ includet("utils/pipeline.jl")
     @testset verbose = true "Jacobian approx" begin
         include("jacobian_approx.jl")
     end
-    @testset verbose = true "Frank-Wolfe" begin
-        include("frank_wolfe.jl")
-    end
     @testset verbose = true "Argmax" begin
         include("argmax.jl")
     end
@@ -30,6 +27,9 @@ includet("utils/pipeline.jl")
     @testset verbose = true "Paths" begin
         include("paths.jl")
     end
+    @testset verbose = true "Frank-Wolfe" begin
+        include("frank_wolfe.jl")
+    end
     @testset verbose = true "Tutorial" begin
         include("tutorial.jl")
     end