Skip to content

feat(aggregation): Add STCH#719

Open
ppraneth wants to merge 2 commits into
SimplexLab:mainfrom
ppraneth:scalarization-3
Open

feat(aggregation): Add STCH#719
ppraneth wants to merge 2 commits into
SimplexLab:mainfrom
ppraneth:scalarization-3

Conversation

@ppraneth
Copy link
Copy Markdown
Contributor

@ppraneth ppraneth commented May 30, 2026

New torchjd.scalarization.STCH, the smooth Tchebycheff scalarization from Smooth Tchebycheff Scalarization for Multi-Objective Optimization .

It returns a differentiable approximation of the weighted, shifted maximum of the values:

$$g_\mu^{\text{STCH}}(v \mid \lambda) = \mu \log \sum_{i=1}^{m} \exp\left(\frac{\lambda_i (f_i - z_i^*)}{\mu}\right)$$

where, following the paper's notation:

  • $f_i$ is the $i$-th input value (the $i$-th objective)
  • $m$ is the number of objectives (the number of elements of the input)
  • $\lambda_i$ is the preference weight for objective $i$ (the weights parameter)
  • $z_i^*$ is the $i$-th component of the ideal point (the reference parameter)
  • $\mu$ is the smoothing parameter (the mu parameter)

As $\mu \to 0$ this recovers the classical (non-differentiable) Tchebycheff $\max_i \lambda_i (f_i - z_i^*)$; larger $\mu$ gives a smoother approximation.

Design decisions

Confirmed with the maintainer before implementation:

  • Paper version, not the reference code. Both the official Xi-L/STCH impl and LibMTL's copy of it are stateful (epoch-based warmup + a running nadir estimate, applied to log(loss / nadir)), which diverges from eq. 9. We use the clean stateless formula the paper proves its theory on.
  • Name: STCH (the paper's acronym).
  • mu is required, no default. The paper tests $\mu \in {0.01, 0.1, 0.5, 1}$ and reports no single value is best across problems, so we force a conscious choice.
  • Full API: mu (required), weights (optional, default uniform on the simplex), reference (optional, default none).

One thing worth a look: the 1/m default and mu

The default weights is uniform on the simplex (1/m), matching the paper. A consequence: the exponent becomes $(f_i - z_i^*) / (m\mu)$, so the effective smoothing temperature is $m\mu$, not $\mu$. In practice this means the meaning of mu is coupled to the number of objectives — more objectives gives a smoother result for the same mu. This is faithful to the paper's simplex convention, but if you'd rather decouple mu from m, the alternative is an all-ones default. Happy to switch if you prefer.

Implementation

def forward(self, values: Tensor, /) -> Tensor:
    # shape checks for weights / reference omitted here
    weights = self.weights if self.weights is not None else torch.full_like(values, 1.0 / values.numel())
    shifted = values if self.reference is None else values - self.reference
    exponents = weights * shifted / self.mu
    return self.mu * torch.logsumexp(exponents.flatten(), dim=-1)
  • Uses torch.logsumexp, so it's numerically stable without manual max-subtraction.
  • No sign precondition on the input (unlike GeometricMean); negative values are fine.
  • mu <= 0 raises in __init__. weights / reference shape mismatches raise at call time (same pattern as Constant).
  • The simplex constraint on user-supplied weights is not enforced (permissive, consistent with Constant).

Files

File Purpose
src/torchjd/scalarization/_stch.py New class
src/torchjd/scalarization/__init__.py Export
docs/source/docs/scalarization/stch.rst Doc page
docs/source/docs/scalarization/index.rst Toctree entry
tests/unit/scalarization/test_stch.py Unit tests
CHANGELOG.md [Unreleased] entry

Test plan

  • uv run pytest tests/unit/scalarization/test_stch.py -W error -v
  • uv run pytest tests/unit -W error (full regression)
  • uv run ruff check && uv run ruff format --check
  • uv run ty check

ppraneth added 2 commits May 29, 2026 20:45
Signed-off-by: ppraneth <pranethparuchuri@gmail.com>
Signed-off-by: ppraneth <pranethparuchuri@gmail.com>
@ppraneth ppraneth requested a review from a team as a code owner May 30, 2026 02:41
@ppraneth ppraneth changed the title add STCH feat(scalarization): add STCH May 30, 2026
@ValerianRey ValerianRey added cc: feat Conventional commit type for new features. package: aggregation labels May 30, 2026
@github-actions github-actions Bot changed the title feat(scalarization): add STCH feat(aggregation): Add STCH May 30, 2026
@github-actions github-actions Bot changed the title feat(scalarization): add STCH feat(aggregation): Add STCH May 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cc: feat Conventional commit type for new features. package: aggregation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants