Add RegularizedGeneric to prediction tests

Author: gdalle

src/InferOpt.jl

@@ -42,7 +42,8 @@ export Interpolation
 
 export DifferentiableFrankWolfe
 
-export shannon_entropy, half_square_norm
+export half_square_norm
+export shannon_entropy, negative_shannon_entropy
 export one_hot_argmax, ranking
 export IsRegularized
 export soft_argmax, sparse_argmax

src/frank_wolfe/differentiable_frank_wolfe.jl

@@ -61,7 +61,7 @@ function frank_wolfe_optimality_conditions(
 end
 
 function ChainRulesCore.rrule(
-    rc::RuleConfig, dfw::DifferentiableFrankWolfe, θ::AbstractVector; fw_kwargs=(;)
+    rc::RuleConfig, dfw::DifferentiableFrankWolfe, θ::AbstractVector{<:Real}; fw_kwargs=(;)
 )
     (; linear_solver) = dfw
 
@@ -94,3 +94,9 @@ function ChainRulesCore.rrule(
 
     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"))
+end

src/regularized/regularized_generic.jl

@@ -24,6 +24,11 @@ struct RegularizedGeneric{ΩF,ΩG,F,G,M,S}
     linear_solver::S
 end
 
+function Base.show(io::IO, regularized::RegularizedGeneric)
+    (; Ω, ∇Ω, maximizer, linear_solver) = regularized
+    return print(io, "RegularizedGeneric($Ω, $∇Ω, $maximizer, $linear_solver)")
+end
+
 function RegularizedGeneric(Ω, ∇Ω, maximizer; linear_solver=gmres)
     f(y, θ) = Ω(y) - dot(θ, y)
     ∇ₓf(y, θ) = ∇Ω(y) - θ
@@ -39,7 +44,7 @@ end
 ## Forward pass
 
 function (regularized::RegularizedGeneric)(
-    θ::AbstractArray; maximizer_kwargs=(;), fw_kwargs=(;)
+    θ::AbstractArray{<:Real}; maximizer_kwargs=(;), fw_kwargs=(;)
 )
     (; f, ∇ₓf, maximizer, linear_solver) = regularized
     lmo = LMOWrapper(maximizer, maximizer_kwargs)
@@ -61,3 +66,13 @@ function ChainRulesCore.rrule(
     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"))
+end

src/regularized/soft_argmax.jl

@@ -6,8 +6,7 @@ Soft argmax activation function `s(z) = (e^zᵢ / ∑ e^zⱼ)ᵢ`.
 Corresponds to regularized prediction on the probability simplex with entropic penalty.
 """
 function soft_argmax(z::AbstractVector{<:Real}; kwargs...)
-    s = exp.(z)
-    s ./= sum(s)
+    s = exp.(z) / sum(exp, z)
     return s
 end
 

test/argmax.jl

@@ -17,54 +17,72 @@ error_function(ŷ, y) = hamming_distance(ŷ, y)
 
 ## Pipelines
 
-pipelines_imitation_θ = [(
-    encoder=encoder_factory(), maximizer=identity, loss=SPOPlusLoss(true_maximizer)
-)]
+pipelines_imitation_θ = [
+    # SPO+
+    (encoder=encoder_factory(), maximizer=identity, loss=SPOPlusLoss(true_maximizer)),
+]
 
 pipelines_imitation_y = [
-    # Fenchel-Young loss (test forward pass)
+    # Structured SVM
     (
         encoder=encoder_factory(),
         maximizer=identity,
-        loss=FenchelYoungLoss(PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=3)),
+        loss=StructuredSVMLoss(ZeroOneBaseLoss()),
+    ),
+    # Perturbed + FYL
+    (
+        encoder=encoder_factory(),
+        maximizer=identity,
+        loss=FenchelYoungLoss(PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=5)),
     ),
     (
         encoder=encoder_factory(),
         maximizer=identity,
         loss=FenchelYoungLoss(PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=5)),
     ),
-    # Other differentiable loss (test backward pass)
+    # Perturbed + other loss
     (
         encoder=encoder_factory(),
-        maximizer=PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=3),
+        maximizer=PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=10),
         loss=Flux.Losses.mse,
     ),
     (
         encoder=encoder_factory(),
-        maximizer=PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=5),
+        maximizer=PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=10),
         loss=Flux.Losses.mse,
     ),
-    # Structured SVM
+    # Explicit regularized + FYL
+    (encoder=encoder_factory(), maximizer=identity, loss=FenchelYoungLoss(sparse_argmax)),
+    (encoder=encoder_factory(), maximizer=identity, loss=FenchelYoungLoss(soft_argmax)),
+    # Explicit regularized + other loss
+    (encoder=encoder_factory(), maximizer=sparse_argmax, loss=Flux.Losses.mse),
+    (encoder=encoder_factory(), maximizer=soft_argmax, loss=Flux.Losses.mse),
+    # Generic regularized + FYL
     (
         encoder=encoder_factory(),
         maximizer=identity,
-        loss=StructuredSVMLoss(ZeroOneBaseLoss()),
+        loss=FenchelYoungLoss(
+            RegularizedGeneric(half_square_norm, identity, true_maximizer)
+        ),
+    ),
+    # Generic regularized + other loss
+    (
+        encoder=encoder_factory(),
+        maximizer=RegularizedGeneric(half_square_norm, identity, true_maximizer),
+        loss=Flux.Losses.mse,
     ),
-    # Regularized prediction: explicit
-    (encoder=encoder_factory(), maximizer=identity, loss=FenchelYoungLoss(sparse_argmax)),
-    (encoder=encoder_factory(), maximizer=identity, loss=FenchelYoungLoss(soft_argmax)),
 ]
 
 pipelines_experience = [
     (
         encoder=encoder_factory(),
         maximizer=identity,
-        loss=cost ∘ PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=3),
+        loss=cost ∘ PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=10),
     ),
     (
         encoder=encoder_factory(),
         maximizer=identity,
-        loss=cost ∘ PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=5),
+        loss=cost ∘ PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=10),
     ),
 ]
 

test/paths.jl

@@ -25,12 +25,13 @@ end
 
 ## Pipelines
 
-pipelines_imitation_θ = [(
-    encoder=encoder_factory(), maximizer=identity, loss=SPOPlusLoss(true_maximizer)
-)]
+pipelines_imitation_θ = [
+    # SPO+
+    (encoder=encoder_factory(), maximizer=identity, loss=SPOPlusLoss(true_maximizer)),
+]
 
 pipelines_imitation_y = [
-    # Fenchel-Young loss (test forward pass)
+    # Perturbed + FYL
     (
         encoder=encoder_factory(),
         maximizer=identity,
@@ -39,21 +40,27 @@ pipelines_imitation_y = [
     (
         encoder=encoder_factory(),
         maximizer=identity,
-        loss=FenchelYoungLoss(
-            PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=10)
-        ),
+        loss=FenchelYoungLoss(PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=5)),
     ),
-    # Other differentiable loss (test backward pass)
+    # Perturbed + other loss
     (
         encoder=encoder_factory(),
-        maximizer=PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=5),
+        maximizer=PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=10),
         loss=Flux.Losses.mse,
     ),
     (
         encoder=encoder_factory(),
         maximizer=PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=10),
         loss=Flux.Losses.mse,
     ),
+    # Generic regularized + FYL
+    (
+        encoder=encoder_factory(),
+        maximizer=identity,
+        loss=FenchelYoungLoss(
+            RegularizedGeneric(half_square_norm, identity, true_maximizer)
+        ),
+    ),
 ]
 
 pipelines_experience = [

test/ranking.jl

@@ -17,37 +17,52 @@ error_function(ŷ, y) = hamming_distance(ŷ, y)
 
 ## Pipelines
 
-pipelines_imitation_θ = [(
-    encoder=encoder_factory(), maximizer=identity, loss=SPOPlusLoss(true_maximizer)
-)]
+pipelines_imitation_θ = [
+    # SPO+
+    (encoder=encoder_factory(), maximizer=identity, loss=SPOPlusLoss(true_maximizer)),
+]
 
 pipelines_imitation_y = [
-    # Fenchel-Young loss (test forward pass)
+    # Interpolation
+    (
+        encoder=encoder_factory(),
+        maximizer=Interpolation(true_maximizer; λ=5.0),
+        loss=Flux.Losses.mse,
+    ),
+    # Perturbed + FYL
     (
         encoder=encoder_factory(),
         maximizer=identity,
-        loss=FenchelYoungLoss(PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=3)),
+        loss=FenchelYoungLoss(PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=5)),
     ),
     (
         encoder=encoder_factory(),
         maximizer=identity,
         loss=FenchelYoungLoss(PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=5)),
     ),
-    # Other differentiable loss (test backward pass)
+    # Perturbed + other loss
     (
         encoder=encoder_factory(),
-        maximizer=PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=3),
+        maximizer=PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=10),
         loss=Flux.Losses.mse,
     ),
     (
         encoder=encoder_factory(),
-        maximizer=PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=5),
+        maximizer=PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=10),
         loss=Flux.Losses.mse,
     ),
-    # Interpolation
+    # Generic regularized + FYL
     (
         encoder=encoder_factory(),
-        maximizer=Interpolation(true_maximizer; λ=5.0),
+        maximizer=identity,
+        loss=FenchelYoungLoss(
+            RegularizedGeneric(half_square_norm, identity, true_maximizer)
+        ),
+    ),
+    # Generic regularized + other loss
+    (
+        encoder=encoder_factory(),
+        maximizer=RegularizedGeneric(half_square_norm, identity, true_maximizer),
         loss=Flux.Losses.mse,
     ),
 ]
@@ -56,14 +71,13 @@ pipelines_experience = [
     (
         encoder=encoder_factory(),
         maximizer=identity,
-        loss=cost ∘ PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=3),
+        loss=cost ∘ PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=10),
     ),
     (
         encoder=encoder_factory(),
         maximizer=identity,
-        loss=cost ∘ PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=5),
+        loss=cost ∘ PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=10),
     ),
-    (encoder=encoder_factory(), maximizer=Interpolation(true_maximizer; λ=5.0), loss=cost),
 ]
 
 ## Dataset generation