Skip to main content

diff_topk

pub fn diff_topk<F>(
    scores: &[F],
    k: usize,
    temperature: F,
) -> OptimizeResult<Vec<F>>
where F: Float + FromPrimitive + Debug + Clone,
Expand description

Differentiable top-k selector via entropy-regularised LP.

Solves the relaxed problem:

max_{p ∈ Δ^n, Σp_i = k}  <scores, p>  -  temperature · H(p)

where H(p) = -Σ p_i log p_i is the entropy regulariser. The solution has the closed form:

p_i = k · softmax(scores / temperature)_i

normalised so that Σp_i = k.

At temperature → 0, p approaches the hard top-k indicator vector.

§Arguments

  • scores – input scores (arbitrary real values).
  • k – number of elements to select (1 ≤ k ≤ n).
  • temperature – regularisation strength (> 0 for differentiable; use a small value like 0.01 for near-hard top-k).

§Returns

Soft indicator vector p in [0,1]^n with Sum(p_i) ~ k.