feat(scalarization): Add cosmos

[ppraneth]

Adds COSMOS from Scalable Pareto Front Approximation for Deep Multi-Objective Learning (Ruchte & Grabocka, ICDM 2021).

It is a stateless scalarizer (loss values only), so it fits Scalarizer.forward(values) directly, like STCH.

Working

COSMOS combines a linear scalarization with a cosine-similarity penalty that pulls the vector of values toward a preference direction:

$$\sum_i r_i L_i ; - ; \lambda , \frac{\sum_i r_i L_i}{\lVert r \rVert , \lVert L \rVert}$$

where $L_i$ is the $i$-th value, $r$ is the preference vector, and $\lambda$ is the penalty coefficient. The subtracted term is $\lambda \cos(r, L)$: minimizing the objective rewards aligning the values with the preference direction, which is what spreads the approximated Pareto front.

Design

• lambda_ (the paper's $\lambda$): required, validated to be non-negative. 0 reduces COSMOS to a plain linear scalarization. The paper uses values from 0.01 to 8 depending on the dataset, with no single best, so there is no default (same choice as STCH's mu). The name follows the SDMGrad precedent for the reserved word.
• weights (the preference $r$): optional, defaults to a uniform preference summing to one, shape-validated at call time. Same convention as STCH.
• No positivity precondition. The cosine term divides by $\lVert L \rVert$, so an all-zero vector of values produces nan; this is documented and locked by a test.

Notes on the formula

I followed the paper's Equation 5 (and the official ruchtem/cosmos implementation), where the linear term is the raw weighted sum $r^\top L$. The libmoon reference normalizes that term by $\lVert r \rVert$, which differs for non-uniform preferences; this is noted in the docstring. Only the scalarization is implemented here. The full COSMOS method also conditions the model on the preference vector by concatenating it to the input, which is a modeling choice left to the user.

Tests

The cosine-aligned zero case (COSMOS(1.0)([1, 1]) == 0), lambda_ = 0 as linear scalarization, custom weights, the full formula with non-uniform weights, structure and gradient flow over all input shapes, permutation invariance with uniform weights, nan on all-zero values, negative lambda_ raising, weights-shape-mismatch raising, and the representations.