google / google/jaxopt

Optimization of complex parameters

Open
#445 2 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
1.1k
Forks
76
Avg merge
2d 21h
Merged PRs (30d)
1

Description

When optimizing a real loss function for a complex variable I encountered some issues with jaxopt 0.7. It seems this use case is less widely used. Are the methods also designed for this use case?

1. I had to use the complex conjugate of the gradient. Is this due to the choice for the complex gradient of jax? Should the jaxopt optimizers take this into account?
2. The 'error' (L2 norm of the gradient), which is used for convergence check is calculated as sum of squares. For complex gradients this should be replaced by the sum over the absolute value, or equivalently, the sum of the squared real and imaginary parts. As a simple fix a patched the l-BGFS and GradientDescent methods. I tested that this change also works with real gradients, and when jitted the extra operation is removed for real gradients.

```python
def tree_l2_norm_patched(tree_x, squared=False):
"""Compute the l2 norm ||tree_x||."""
squared_tree = jaxopt.tree_util.tree_map(lambda x: jnp.square(x.real) + jnp.square(x.imag), tree_x)
sqnorm = jaxopt.tree_util.tree_sum(squared_tree)
if squared:
return sqnorm
else:
return jnp.sqrt(sqnorm)

import jaxopt._src
jaxopt._src.lbfgs.tree_l2_norm = tree_l2_norm_patched
jaxopt._src.proximal_gradient.tree_l2_norm = tree_l2_norm_patched
```

Gregor

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.