Support objective functions with auxiliary variables #444

bagibence · 2025-11-21T14:57:06Z

Problem stated in #375:

Optimization libraries support objective functions that return not just a scalar function value, but also auxiliary data. These are mostly used for logging, diagnostics, and debugging.
Currently, the objective of GLMs does not have any aux, and throughout the codebase (e.g. in the AbstractSolver interface) nemos assumes that objective functions return a scalar only.

As other models might make use of aux, this PR prepares the solver interface and current models to deal with that.

Main changes:

(Prox-)SVRG can now handle objectives with aux. Following JAXopt solvers, aux is saved in the solver state. This saved aux does not come from the last evaluation of the objective or its gradient -- which is done on a minibatch -- but is the result of evaluating the gradient on the full data at the last reference point.
BaseRegressor has a class attribute called has_aux that is passed to the solver on instantiation. Models whose objective returns aux have to overwrite this.
Accordingly, has_aux is now a required argument of AbstractSolver.__init__.
Adapt tests to these changes.

Remaining questions, tasks:

In order to avoid braking existing code, GLM.update doesn't return aux but saves it in self.aux. Should this be named self.aux_ instead? Or returned?
GLM.update performed the update, then estimated the scale using the previous parameters. Was that intended? I changed it to use the new parameters.
has_aux is the only class attribute of BaseRegressor. Is that okay or should it be stored somewhere else? In any case, I will add a note about it to the developer notes.

Fixes #375

Child classes whose objective function returns aux will have to set this to true.

There was a potential bug where at the end of GLM.update self.scale_ was estimated using the old params.

Update the test for it as well.

BalzaniEdoardo · 2025-11-25T15:53:57Z

Remaining questions, tasks:

* In order to avoid braking existing code, GLM.update doesn't return aux but saves it in `self.aux`. Should this be named `self.aux_` instead? Or returned?

I would call it self.aux_

* `GLM.update` performed the update, then estimated the scale using the previous parameters. Was that intended? I changed it to use the new parameters.

good catch

* `has_aux` is the only class attribute of `BaseRegressor`. Is that okay or should it be stored somewhere else? In any case, I will add a note about it to the developer notes.

I think that's ok.

BalzaniEdoardo

I left a few comments but all minor stuff

src/nemos/glm/glm.py

BalzaniEdoardo · 2025-11-25T15:58:27Z

src/nemos/glm/glm.py

        # the output of loss. I believe it's the output of
        # solver.l2_optimality_error
        self.solver_state_ = state
+        # TODO: Should this be part of fit-state, so called aux_?


probably also part of the fit state so that people running a loop externally, can concatenate the aux across iterations? if there is some metric that users may want to track, it would make it easier

Do you mean solver state? By fit-state I meant what is extracted by BaseRegressor._get_fit_state

BalzaniEdoardo · 2025-11-25T16:02:25Z

src/nemos/solvers/_optimistix_solvers.py

+    f_struct is "the shape+dtype of the output of `fn`".
+    aux_struct is the same for the returned aux.
+    """
+    y0 = jax.tree_util.tree_map(optx._misc.inexact_asarray, y0)


what does this inexact_asarray do? and why are we using a private function? is there a public equivalent?

BalzaniEdoardo · 2025-11-25T16:12:44Z

src/nemos/solvers/_optimistix_solvers.py

+    y0 = jax.tree_util.tree_map(optx._misc.inexact_asarray, y0)
+    if not has_aux:
+        fn = optx._misc.NoneAux(fn)  # pyright: ignore
+    fn = optx._misc.OutAsArray(fn)


These _misc functions are tiny wrappers - since they're internal utilities, we should consider porting them directly instead of relying on optx._misc for maintainability.

Specifically:

The wrapper module has a __call__ that returns (func(x), None)

inexact_asarray is just jnp.asarray(x) which converts numeric scalars to arrays at default float precision (per JAX docs: "all numeric scalar types with a (potentially) inexact representation")

Since these are so small, it would be more maintainable to inline them rather than depend on private Optimistix APIs.

probaly OutAsArray is also another small one

I see your point. I'm not sure yet. They're all small, but they call a bunch of other ones that would have to be ported as well, and it adds up.
For now I removed the dependency and made my own little wrappers for the absolutely necessary ones.
I'll look into if and why the others are need.

BalzaniEdoardo · 2025-11-25T16:16:45Z

src/nemos/solvers/_optimistix_solvers.py

+            self.fun = lambda params, args: loss_fn(params, *args)[0]
+        else:
+            self.fun = lambda params, args: loss_fn(params, *args)
+            self.fun_with_aux = lambda params, args: (loss_fn(params, *args), None)


you can directly use optx._misc.NoneAux or its port if we have it. this is the same call used inside the function that returns the srcut of f and aux, so at this point we can have both matching

Using the locally defined _wrap_aux now.

BalzaniEdoardo · 2025-11-25T16:19:40Z

src/nemos/solvers/_svrg.py

-                    prev_reference_point, *args
-                )
+                full_grad_at_reference_point=full_grad,
+                aux=new_aux,


nice, so we have the aux as part of the state. is that true for every solver already?

This is true for all jaxopt and optmistix solvers we have

I don't think so. Optax doesn't support auxiliary variables and so optimistix.OptaxMinimiser, and with that OptimistixOptaxGradientDescent and OptimistixOptaxLBFGS, don't have it in the state.

macari216 · 2025-11-27T03:18:05Z

src/nemos/solvers/_jaxopt_solvers.py

    def run(self, init_params: Params, *args: Any) -> JaxoptStepResult:
-        return self._solver.run(init_params, *self.hyperparams_prox, *args)
+        params, state = self._solver.run(init_params, *self.hyperparams_prox, *args)
+        return (params, state, state.aux)


Just curious, is there a reason for returning aux explicitly in solver.run and solver.update? As opposed to returning just params and state (following the internal _solver API) and then setting self.aux_ = opt_state.aux in glm fit and update.

That would work great if we knew that opt_state.aux exists for sure, but not every solver is guaranteed to store it in the state.
It's also consistent with how step works in Optimistix.

bagibence added 23 commits November 21, 2025 11:07

Make (Prox)SVRG handle objective functions with aux

39fa23b

Update solver interface to use aux

f23dcd3

Fix bug in SVRG's aux handling

8511872

Fix bug in Optimistix solvers' aux handling

b93c034

Comments, format

0725757

Make jaxopt solvers accept has_aux

c278edd

Add tests for aux handling

0911d9f

tiny check in test

98b872f

Comments

ae28e04

Add has_aux attribute in BaseRegressor and pass it to the solver

ec39dd4

Child classes whose objective function returns aux will have to set this to true.

Fix handling has_aux in jaxopt solvers

268ef39

Update GLM to handle aux returned by solver. Also bug fix?

9c32cd1

There was a potential bug where at the end of GLM.update self.scale_ was estimated using the old params.

Update test_regularizer to handle aux returned by solver

66ff33c

Save aux in GLM

8d5d3a2

Fix docstring

6d3dd2d

Typing in GLM

ce586d6

Update how f_struct and aux_struct are inferred in OptimistixAdapter

ffbff34

Update the test for it as well.

Add todo

be8455e

Formatting

688def5

Update test, remove comment

576299d

Require has_aux in AbstractSolver.__init__

8f95986

Rename BaseRegressor.has_aux to _has_aux

7826108

Docs

1672f9d

bagibence marked this pull request as ready for review November 24, 2025 09:42

bagibence requested a review from BalzaniEdoardo as a code owner November 24, 2025 09:42

BalzaniEdoardo requested a review from macari216 November 25, 2025 15:51

BalzaniEdoardo requested changes Nov 25, 2025

View reviewed changes

macari216 reviewed Nov 27, 2025

View reviewed changes

Remove todo

8a2b0b5

bagibence added 2 commits November 27, 2025 15:23

Rename BaseRegressor/GLM.aux -> aux_ and adjust tests

9406675

Stop relying on optimistix._misc

54c7f7d

Support objective functions with auxiliary variables #444

Are you sure you want to change the base?

Support objective functions with auxiliary variables #444

Uh oh!

Conversation

bagibence commented Nov 21, 2025

Problem stated in #375:

Main changes:

Remaining questions, tasks:

Uh oh!

BalzaniEdoardo commented Nov 25, 2025

Remaining questions, tasks:

Uh oh!

BalzaniEdoardo left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants