Reference

Base type of training objectives.

(obj::AdversarialObjective)(
    yhat,
    y,
    energy_differential::AbstractVector{<:AbstractFloat}=[0.0f0],
    regularization::AbstractVector{<:AbstractFloat}=[0.0f0],
    adversarial_loss::Union{AbstractFloat,AbstractVector{<:AbstractFloat}}=0.0f0;
    agg=mean,
    kwrgs...,
)

If the adversarial_loss has been computed already, obj::AdversarialObjective can be called directly on predictions yhat and labels y. The different loss components are then added together with a weighting vector lambda.

AdversarialObjective <: AbstractObjective

The AdversarialObjective is a concrete implementation of the AbstractObjective abstract type that optimizes for:

1. Standard classification objective (the discriminative task).
2. Adversarial classification objective on the counterfactuals (the explainability task).

AdversarialObjective(;
    class_loss::Function=Flux.Losses.logitcrossentropy,
    lambda::Vector{<:AbstractFloat}=[1.0, 0.25]
)

Outer constructor for the AdversarialObjective type.

(obj::EnergyDifferentialObjective)(
    yhat,
    y,
    energy_differential::AbstractVector{<:AbstractFloat}=[0.0f0],
    regularization::AbstractVector{<:AbstractFloat}=[0.0f0],
    adversarial_loss::Union{AbstractFloat,AbstractVector{<:AbstractFloat}}=0.0f0;
    agg=mean,
    kwrgs...,
)

If the energy_differential and regularization have been computed already, obj::EnergyDifferentialObjective can be called directly on predictions yhat and labels y. The different loss components are then added together with a weighting vector lambda.

EnergyDifferentialObjective <: AbstractObjective

The EnergyDifferentialObjective is a concrete implementation of the AbstractObjective abstract type that optimizes for:

1. Standard classification objective (the discriminative task)
2. Energy differential between counterfactuals and observed data (the explainability task).

EnergyDifferentialObjective(;
    class_loss::Function=Flux.Losses.logitcrossentropy, 
    lambda::Vector{<:AbstractFloat}=[1.0, 0.5, 0.1]
)

Outer constructor for the EnergyDifferentialObjective type.

(obj::FullObjective)(
    yhat,
    y,
    energy_differential::AbstractVector{<:AbstractFloat}=[0.0f0],
    regularization::AbstractVector{<:AbstractFloat}=[0.0f0],
adversarial_loss::Union{AbstractFloat,AbstractVector{<:AbstractFloat}}=[0.0f0];
    agg=mean,
    kwrgs...,
)

If the adversarial_loss has been computed already, obj::FullObjective can be called directly on predictions yhat and labels y. The different loss components are then added together with a weighting vector lambda.

FullObjective <: AbstractObjective

The FullObjective is a concrete implementation of the AbstractObjective abstract type that optimizes for all three tasks:

1. Standard classification objective (the discriminative task)
2. Energy differential between counterfactuals and observed data (the explainability task).
3. Adversarial classification objective on the counterfactuals (the explainability task).

FullObjective(;
    class_loss::Function=Flux.Losses.logitcrossentropy,
    energy_differential::PenaltyOrFun=EnergyDifferential(),
    lambda::Vector{<:AbstractFloat}=[1.0, 0.5, 0.1, 0.25]
)

Outer constructor for the FullObjective type.

(obj::VanillaObjective)(
    yhat,
    y,
    energy_differential::AbstractVector{<:AbstractFloat}=[0.0f0],
    regularization::AbstractVector{<:AbstractFloat}=[0.0f0],
    adversarial_loss::Union{AbstractFloat,AbstractVector{<:AbstractFloat}}=0.0f0;
    agg=mean,
    kwrgs...,
)

obj::VanillaObjective can be called directly on predictions yhat and labels y.

VanillaObjective <: AbstractObjective

The VanillaObjective is a concrete implementation of the AbstractObjective abstract type that optimizes for:

1. Standard classification objective (the discriminative task).

VanillaObjective(class_loss, lambda; needs_ce=false)

Outer constructor to allow passing just class_loss and lambda as positional arguments.

VanillaObjective(;
    class_loss::Function=Flux.Losses.logitcrossentropy,
    lambda::Vector{<:AbstractFloat}=[1.0],
    needs_ce::Bool=false,
)

Outer constructor for the VanillaObjective type.

accuracy(model, train_set; device=identity)

Compute classification accuracy over a DataLoader. Evaluates the model in test (eval) mode so BatchNorm uses running rather than batch statistics, and accumulates match counts on the device, syncing to the host once per call. The device keyword moves each batch to the device before the forward pass.

adv_loss(
    model, counterfactual, perturbations, targets; epsilon=2.0, p::Real=Inf, validities=nothing
)

Adversarial loss function.

generate!(
    model,
    data,
    generator::AbstractGenerator;
    nsamples::Union{Int,Nothing}=nothing,
    convergence=Convergence.MaxIterConvergence(),
    parallelizer=nothing,
    input_encoder=nothing,
    verbose=1,
    domain=nothing,
)

This function generates counterfactual explanations for the whole dataset data or a subset thereof (nsamples). It is supposed to be used outside of the mini-batch training loop.

get_global_ae_criterium()

Get the global AE criterium.

get_last_valid_ae(ce::CounterfactualExplanation)

A callback function used to store the last counterfactual that is also a valid adversarial example based on the global AE criterium (see get_global_ae_criterium).

implausibility(model, counterfactual, samples, targets)

Compute the implausibility (contrastive divergence) of the counterfactuals (counterfactual) with respect to samples in the target class. This is computed as the difference between negative logits indexed at the target class for the samples and the counterfactual.

implausibility_and_reg_loss(model, counterfactual, samples, targets)

Computes both implausibility and reg_loss in a single pass, sharing the forward passes through model for samples and counterfactual. Returns (implaus, regs) — the same values that implausibility(...) and reg_loss(...) would return separately.

This avoids redundant forward passes when both losses are needed (e.g. inside the training loop's gradient tape).

implausibility_and_reg_loss_from_logits(logits_cf, logits_nb, targets)

Like implausibility_and_reg_loss but accepts precomputed logits (logits_cf for the counterfactuals and logits_nb for the neighbours) instead of calling model internally. Used by the training loop when fuse_cf_forwards=true to share a single concatenated forward pass between the implausibility/regularization computation and the adversarial loss.

The algebra matches implausibility_and_reg_loss exactly.

infer_domain_constraints(X::AbstractArray; nstd=3)

Automatically infers reasonable domain constraints for the counterfactuals.

isvalid(ce, model, data)

Checks if the label has been flipped.

protect_immutable!(
    samples::AbstractArray,
    counterfactuals::AbstractArray,
    mutability::Union{Nothing,AbstractArray},
)

Protects immutable features from the contrastive divergence penalty.

reg_loss(model, counterfactual, samples, targets)

Compute the regularization loss for the contrastice divergence.

setup_counterfactual_search(
    data,
    model,
    domain,
    input_encoder,
    mutability,
    nneighbours::Int64,
    nsamples::Union{Nothing,Int64},
)

Sets up the counterfactual search.

unwrap(train_set; labels=nothing)

Unwraps the data from a Flux.DataLoader or zip iterator. The output variables is assumed to be categorical. If no labels are provided, then 1 to n is used where n is the number of classes.

NativeGenerator

Lightweight generator for batched counterfactual search. Holds only the fields needed by generate_counterfactuals!: penalty weights λ and a Flux optimiser opt.

This avoids the Flux 0.16 incompatibility in CE.jl's GradientBasedGenerator (whose opt field is typed against the old Flux.Optimise.AbstractOptimiser).

Keyword fields

• λ::Vector{Float32} — penalty weights [λ₁, λ₂] for the L1 distance penalty and the energy constraint (default [0.1f0, 1.0f0]).
• opt — a Flux optimiser (default Flux.Descent(0.1f0)).

batched_apply_domain_constraints!(X′::AbstractMatrix, data::CounterfactualData; device=identity)

Clamps each feature row of X′ to the domain bounds stored in data.domain. Uses vectorized broadcasting for GPU compatibility. The device keyword moves the bounds arrays to the compute device before broadcasting, preventing non-bitstype CPU array capture in GPU kernels.

This is a convenience wrapper that builds the bounds on the fly and delegates to batched_apply_domain_constraints!(X′, bounds). For repeated calls with the same data, precompute bounds with prepare_domain_bounds and call the bounds-accepting signature directly.

batched_apply_domain_constraints!(X′::AbstractMatrix, bounds)

Clamp each feature row of X′ to the precomputed domain bounds in-place. bounds is a 2-tuple (lb, ub) as returned by prepare_domain_bounds, or nothing (no-op).

Uses vectorized broadcasting for GPU compatibility.

batched_apply_mutability!(ΔX::AbstractMatrix, masks)

Apply precomputed mutability masks to a batched gradient update in-place. masks is a 3-tuple (none_mask, inc_mask, dec_mask) as returned by prepare_mutability_masks, or nothing (no-op).

Zeros out gradient components along immutable feature directions in-place. Uses vectorized broadcasting for GPU compatibility.

batched_apply_mutability!(ΔX::AbstractMatrix, mutability; device=identity)

Zeros out gradient components along immutable feature directions in-place. mutability is a vector of Symbols (:both, :none, :increase, :decrease), one per feature row. Uses vectorized broadcasting for GPU compatibility. The device keyword moves the mask arrays to the compute device before broadcasting, preventing non-bitstype CPU array capture in GPU kernels.

This is a convenience wrapper that builds the masks on the fly and delegates to batched_apply_mutability!(ΔX, masks). For repeated calls with the same mutability, precompute masks with prepare_mutability_masks and call the mask-accepting signature directly.

batched_energy(model, X′::AbstractMatrix, target_idx::AbstractVector)

Returns a length-N vector of negative logits at the target class for each sample in X′. Calls model(X′) internally and delegates to batched_energy_from_logits.

batched_energy_from_logits(logits::AbstractMatrix, target_idx::AbstractVector{Int})

Returns a length-N vector of negative logits at the target class for each sample, given the C×N logits matrix. Uses linear indexing for GPU compatibility.

This is the core indexing logic extracted from batched_energy so that callers who already have logits = model(X′) can avoid a redundant forward pass.

check_batched_convergence(probs, target_idx, iter, maxiter, threshold)

Returns a BitVector of length N indicating which samples have converged. Convergence is reached when the target-class probability ≥ threshold or iter ≥ maxiter.

find_neighbours(X, y, targets, y_levels; nneighbours=1, rng=Random.default_rng())

For each counterfactual with target targets[i], samples a random training point that has label targets[i]. Returns a D×N matrix (one neighbour per counterfactual column).

Uses a precomputed class→indices dictionary to avoid per-sample findall calls. The nneighbours keyword is accepted for API compatibility but currently only one neighbour per sample is returned.

generate_counterfactuals!(
    model,
    X::AbstractMatrix,
    targets::Vector{Int},
    data::CounterfactualData,
    generator::NativeGenerator;
    maxiter = 30,
    decision_threshold = 0.75f0,
    decay = 0.9f0,
    reg_strength = 1.0f-3,
    epsilon = 0.3f0,
    p = Inf,
    device = identity,
    cf_batchsize = 128,
)

Generates counterfactual explanations for a batch of N factuals in a fully batched (GPU-compatible) fashion.

Arguments

• model: A Flux model (or any callable that accepts a D×N matrix and returns a C×N matrix of logits).
• X: D×N matrix of factuals.
• targets: Length-N vector of target class indices (1-based integers).
• data: A CounterfactualData object providing domain bounds and mutability.
• generator: A NativeGenerator.

Keyword arguments

• maxiter: Maximum number of search iterations.
• decision_threshold: Target probability threshold for convergence.
• decay: Decay rate for the polynomial decay schedule.
• reg_strength: Regularization strength for the energy penalty.
• epsilon: Norm bound for tracking valid adversarial examples.
• p: Norm order for the adversarial example bound (default Inf).
• device: Function to move data to the compute device (identity for CPU, Flux.gpu for GPU). Factuals and one-hot targets are moved to the device before the search loop.
• cf_batchsize: Mini-batch size for the counterfactual search forward/backward passes. Controls peak GPU memory: the search processes cf_batchsize samples at a time through the model. Default 128. Set to a larger value for GPUs with more memory, or smaller for memory-constrained GPUs. When cf_batchsize >= nsamples, no chunking occurs.

Returns

• counterfactuals::AbstractMatrix: The final counterfactuals (D×N).
• last_valid_ae::AbstractMatrix: Last valid adversarial examples (D×N).
• converged_mask::BitVector: Per-sample convergence flag.
• maxiter::Int: The maximum number of iterations used.

generate_native!(
    model, train_set, generator::NativeGenerator;
    nsamples=nothing, nneighbours=1, domain=nothing, mutability=nothing,
    maxiter=30, decision_threshold=0.75f0, decay=0.9f0,
    reg_strength=1.0f-3, epsilon=0.3f0, p=Inf, verbose=1, device=identity,
    cf_batchsize=128,
    cached_X=nothing, cached_y_raw=nothing, cached_data=nothing,
)

Top-level counterfactual generation for the native branch. Generates counterfactuals for a subset of the training data, finds neighbours, applies mutability protection, and partitions results into a data loader aligned with train_set batches.

The device keyword (a function: identity for CPU, Flux.gpu for GPU) moves the subsampled factuals to the device for model calls and the counterfactual search. Results stay on the compute device; the returned data loader contains device arrays ready for use in the training loop.

Returns (dl, percent_valid, nothing) — same interface as the old generate!().

The factual-prediction loop and the counterfactual search run the model in test (eval) mode (train mode is restored afterwards). This keeps BatchNorm running statistics clean of adversarial inputs and makes the CF search deterministic w.r.t. chunk size. For BN-free models this is a no-op.

Keyword arguments

• cf_batchsize: Mini-batch size for the counterfactual search forward/backward passes. Controls peak GPU memory: the search processes cf_batchsize samples at a time through the model. Default 128. Set to a larger value for GPUs with more memory, or smaller for memory-constrained GPUs. When cf_batchsize >= nsamples, no chunking occurs.

Cached keyword arguments

• cached_X: Pre-unwrapped feature matrix (CPU). When provided (by counterfactual_training), avoids calling unwrap(train_set) every epoch. When nothing (default, e.g. standalone calls), unwrap is called on the fly.
• cached_y_raw: Pre-unwrapped label vector (CPU). Paired with cached_X.
• cached_data: Pre-built CounterfactualData object. When provided, the domain and mutability keyword arguments are ignored (they are already baked into cached_data). When nothing (default), CounterfactualData is constructed from X, y_raw, domain, and mutability.

generator_loss(gen, model, X′, X, targets_onehot, target_idx, iter,
               reg_strength, decay, maxiter)

Batched version of the ECCo objective. Operates on D×N matrices.

generator_loss_from_logits(gen, logits, X′, X, targets_onehot, target_idx, iter,
                           reg_strength, decay, maxiter)

Like generator_loss but accepts precomputed logits = model(X′) instead of calling model internally. Used by generate_counterfactuals! to share the forward pass between the gradient computation and the convergence check.

prepare_domain_bounds(domain, D; device=identity)

Precompute the lower/upper bound arrays used by batched_apply_domain_constraints!.

Returns (lb, ub), each a D×1 array on the compute device, or nothing if domain is nothing.

prepare_mutability_masks(mutability, D; device=identity)

Precompute the per-direction boolean masks used by batched_apply_mutability!.

Returns (none_mask, inc_mask, dec_mask), each a D×1 boolean array on the compute device, or nothing if mutability is nothing.

protect_immutable!(neighbours, counterfactuals, masks)

Protects immutable features by setting neighbour values according to mutability directions. For each feature row:

• :both → keep neighbour value
• :none → use counterfactual value (no change)
• :increase → max(counterfactual, neighbour)
• :decrease → min(counterfactual, neighbour)

Accepts precomputed masks (a 3-tuple of D×1 boolean arrays as returned by prepare_mutability_masks), or nothing (no-op). The mask-based signature uses broadcast ifelse. with D×1 masks against D×N arrays, making it GPU-safe.

The old signature accepting a Vector{Symbol} is kept as a back-compat wrapper that builds masks on the fly (CPU-only).

split_obs(indices, n)

Split indices into n roughly equal groups. Returns a vector of vectors.

track_adversarial_examples!(last_valid_ae, X, X′, epsilon, p, perturbations, norms)

In-place version that writes the perturbation matrix into perturbations and the per-sample norms into norms (both preallocated, same shape/type as needed). Updates last_valid_ae in-place: for samples whose perturbation norm (in p-norm) is ≤ epsilon, the current counterfactual is stored.

track_adversarial_examples!(last_valid_ae, X, X′, epsilon, p)

Convenience wrapper that allocates perturbations and norms on the fly. For repeated calls (e.g. inside the counterfactual search loop), prefer the 6-argument signature with preallocated buffers.

counterfactual_training(
    loss::AbstractObjective,
    model,
    generator::NativeGenerator,
    train_set,
    opt_state;
    device = identity,
    val_set = nothing,
    nepochs = 100,
    burnin = 0.0f0,
    nce = nothing,
    nneighbours = 100,
    domain = nothing,
    mutability = nothing,
    maxiter = 30,
    decision_threshold = 0.75f0,
    decay = 0.9f0,
    reg_strength = 1.0f-3,
    epsilon = 0.3f0,
    p = Inf,
    verbose = 1,
    checkpoint_dir = nothing,
    callback = nothing,
    cf_batchsize = 128,
    accuracy_every::Real = Inf,
    fuse_cf_forwards::Bool = false,
)

Native GPU-compatible counterfactual training. Dispatches here when the generator is a NativeGenerator.

The device keyword is a function: identity (CPU), Flux.gpu (CUDA), or AMDGPU.gpu (AMDGPU). The model is moved to the device; training data should already be on the device (user moves it before constructing the DataLoader).

Keyword arguments

• cf_batchsize: Mini-batch size for the counterfactual search forward/backward passes. Controls peak GPU memory: the search processes cf_batchsize samples at a time through the model. Default 128. Set to a larger value for GPUs with more memory, or smaller for memory-constrained GPUs. When cf_batchsize >= nsamples, no chunking occurs.
• accuracy_every: Compute training (and validation) accuracy only every accuracy_every epochs. When epoch % accuracy_every != 0, the logged acc and acc_val fields are nothing. Default Inf (accuracy is never computed unless explicitly requested). Set to 1 for every epoch, or a larger value (e.g. 10) to reduce per-epoch wall-clock time for large models and datasets.
• fuse_cf_forwards: When true, the three counterfactual forward passes in the training loop (perturbed_input, neighbours, advexms) are fused into a single concatenated forward pass, reducing kernel-launch overhead for launch-bound workloads. Caveat: fusing changes BatchNorm batch statistics (stats are computed over the concatenated mini-batch rather than each tensor at its native width). For BN-free models results are identical; for BN models results differ slightly. Off by default.