Oracles

Oracles

ImplicitOracle

Bases: Oracle

Oracle using implicit differentiation through the lower-level optimum.

For given group weights lambda, the lower-level model is fit with these weights. The gradient uses the implicit function theorem: it forms a mixed Hessian of the lower-level loss (weighted sum over groups) and solves linear systems to obtain dw*/dlambda. The final gradient is (dw*/dlambda)^T @ d(metric)/dw evaluated at w*.

Notes

• Requires the model to provide a differentiable scalar loss model._loss(w, X, y).
• Uses JAX grad and hessian; the solve uses jax.numpy.linalg.solve.

Methods:

• solve_lower –

  Fit the model with group_weights and return the resulting coefficients.

• w_star_grad –

  Compute dw*/dgroup_weights at the current solution.

fun(group_weights)

Evaluate the upper-level metric at the fitted lower-level solution.

Parameters:

• group_weights ((ndarray, shape(n_groups))) –

  Group weights.

Returns:

• ArrayLike –

  Metric value metric.fun(w*, dset, train_test).

grad(group_weights)

Gradient of the upper-level metric w.r.t. group weights.

Parameters:

• group_weights ((ndarray, shape(n_groups))) –

  Group weights.

Returns:

• ArrayLike –

  Gradient vector of shape (n_groups,).

solve_lower(group_weights: jnp.ndarray) -> jnp.ndarray

Solve the lower-level problem for fixed group weights.

Parameters:

• group_weights ((ndarray, shape(n_groups))) –

  Group weights used by the model during training.

Returns:

• ndarray –

  Coefficient vector at the fitted solution.

w_star_grad(group_weights: jnp.ndarray) -> jnp.ndarray

Jacobian of the lower-level optimum w* w.r.t. group weights.

Parameters:

• group_weights ((ndarray, shape(n_groups))) –

  Group weights.

Returns:

• (ndarray, shape(n_params, n_groups)) –

  Matrix whose column i is dw*/dlambda_i.

Oracle

Base oracle for the bilevel BADR objective.

An oracle takes group weights lambda and returns the upper-level objective value and gradient w.r.t. lambda. Subclasses differ in how they handle the lower-level optimization over model parameters.

Parameters:

• dset (Dataset) –

  Dataset providing splits and group indices.

• model (Model) –

  Model used for the lower-level problem.

• metric (FairnessMetric) –

  Upper-level metric evaluated at the lower-level solution.

• train_test ((train, test), default: "train" ) –

  Split used to form groups and evaluate the metric.

Attributes:

• X, y –

  Feature matrix and targets for the chosen split.

• groups (list[ndarray]) –

  Group indices for the chosen split.

• n_groups (int) –

  Number of groups.

fun(group_weights: jnp.ndarray) -> ArrayLike

Upper-level objective value at group_weights.

Parameters:

• group_weights ((ndarray, shape(n_groups))) –

  Group weights (often denoted lambda).

Returns:

• ArrayLike –

  Scalar objective value.

Raises:

• NotImplementedError –

  If not implemented.

grad(group_weights: jnp.ndarray) -> ArrayLike

Gradient of :meth:fun w.r.t. group_weights.

Parameters:

• group_weights ((ndarray, shape(n_groups))) –

  Group weights.

Returns:

• ArrayLike –

  Gradient vector of shape (n_groups,).

Raises:

• NotImplementedError –

  If not implemented.

StochasticOracle

Bases: Oracle

Oracle based on stochastic minibatches from each group.

This oracle keeps per-group views of the data. It provides utilities to: - sample balanced minibatches across groups, - compute a stochastic estimate of the weighted lower-level objective f_hat(w, lambda) = sum_s lambda[s] * loss_s(w; batch_s), - compute gradients and Hessian-vector products w.r.t. w using JAX AD.

Notes

• Sampling uses a JAX PRNG key stored on the instance.
• Per-group losses are computed via model._loss(w, X_batch, y_batch).
• hvp_* methods use jax.jvp to avoid forming full Hessians.

Attributes:

• X_list, y_list (list[ndarray]) –

  Per-group feature/target arrays for the chosen split.

• S (int) –

  Number of groups.

• d (int) –

  Number of features.

• key (Array) –

  PRNG key used for sampling.

f_hat(w: jnp.ndarray, lam: jnp.ndarray, batch: _Batch) -> jnp.ndarray

Stochastic weighted lower-level objective on a batch.

Parameters:

• w (ndarray) –

  Model parameters.

• lam ((ndarray, shape(S))) –

  Group weights.

• batch (_Batch) –

  Per-group minibatch.

Returns:

• ndarray –

  Scalar value lam @ [loss_s(w; batch_s)].

grad_lower_groups(w: jnp.ndarray, batch_or_size: Union[_Batch, int], key: Optional[jax.Array] = None) -> jnp.ndarray

Per-group gradients of the group losses at w.

Parameters:

• w (ndarray) –

  Model parameters.

• batch_or_size (_Batch or int) –

  Either an explicit batch, or a batch size to sample.

• key (Array, default: None ) –

  PRNG key used when sampling.

Returns:

• (ndarray, shape(S, n_params)) –

  Stack of gradients, one per group.

grad_upper(w: jnp.ndarray, lam: jnp.ndarray)

Upper-level gradients at w.

Computes the gradient of the fairness metric w.r.t. w. The gradient w.r.t. lam is zero here because the metric is evaluated only on w (no direct dependence on lam in this implementation).

Parameters:

• w (ndarray) –

  Model parameters.

• lam (ndarray) –

  Group weights.

Returns:

• (ndarray, ndarray) –

  (d metric / d w, d metric / d lam).

grad_w_f_hat(w: jnp.ndarray, lam: jnp.ndarray, batch: _Batch) -> jnp.ndarray

Gradient of :meth:f_hat w.r.t. w.

Parameters:

• w (ndarray) –

  Model parameters.

• lam ((ndarray, shape(S))) –

  Group weights.

• batch (_Batch) –

  Per-group minibatch.

Returns:

• ndarray –

  Gradient vector.

hvp_lower_group(w: jnp.ndarray, v: jnp.ndarray, group_idx: int, batch_size: int = 32, key: Optional[jax.Array] = None) -> jnp.ndarray

Hessian-vector product for a single group's loss.

Parameters:

• w (ndarray) –

  Model parameters.

• v (ndarray) –

  Vector to multiply by the Hessian.

• group_idx (int) –

  Group index.

• batch_size (int, default: 32 ) –

  Batch size for the group minibatch.

• key (Array, default: None ) –

  PRNG key used when sampling.

Returns:

• ndarray –

  H_w(loss_group) @ v.

hvp_w_f_hat(w: jnp.ndarray, lam: jnp.ndarray, v: jnp.ndarray, batch: _Batch) -> jnp.ndarray

Hessian-vector product of :meth:f_hat w.r.t. w.

Parameters:

• w (ndarray) –

  Model parameters.

• lam ((ndarray, shape(S))) –

  Group weights.

• v (ndarray) –

  Vector to multiply by the Hessian.

• batch (_Batch) –

  Per-group minibatch.

Returns:

• ndarray –

  H_w(f_hat) @ v.

jt_v_lambda_of_grad_w_f_hat(w: jnp.ndarray, v: jnp.ndarray, batch: _Batch) -> jnp.ndarray

Compute J^T v where J = d/dlambda (grad_w f_hat).

For f_hat(w, lambda) = sum_s lambda[s] * loss_s(w), grad_w f_hat = sum_s lambda[s] * grad_w loss_s and the Jacobian w.r.t. lambda has columns grad_w loss_s. This method returns the vector of dot-products [<grad_w loss_s, v>] over groups.

Parameters:

• w (ndarray) –

  Model parameters.

• v (ndarray) –

  Vector in parameter space.

• batch (_Batch) –

  Per-group minibatch.

Returns:

• (ndarray, shape(S)) –

  [dot(grad_w loss_s, v)]_s.

sample_batch(key: jax.Array, batch_size: int = 32) -> Tuple[_Batch, jax.Array]

Sample a (roughly) balanced minibatch across groups.

The batch is split as evenly as possible: each group gets batch_size // S samples, and the remainder is distributed to the first groups.

Parameters:

• key (Array) –

  PRNG key.

• batch_size (int, default: 32 ) –

  Total batch size across all groups.

Returns:

• _Batch –

  Batch containing per-group X/y slices and per-group counts.

• Array –

  Updated PRNG key.

sample_group_batch(key: jax.Array, group_idx: int, batch_size: int = 32) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], jax.Array]

Sample a minibatch from a single group.

Parameters:

• key (Array) –

  PRNG key.

• group_idx (int) –

  Group index in [0, S).

• batch_size (int, default: 32 ) –

  Maximum number of samples to draw.

Returns:

• (ndarray, ndarray) –

  (X_batch, y_batch).

• Array –

  Updated PRNG key.