google / google/jaxopt

Stopping condition 'madsen-nielsen' incorrect

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

Description

The documentation says a different thing than code. Specifically, the `-` is inconsistent with the `+` in docstring at this part `(tree_l2_norm(params) - self.xtol)`.

Docstring says:
```
the convergence is achieved once the
coeff update satisfies ``||dcoeffs||_2 <= xtol * (||coeffs||_2 + xtol) `` or
the gradient satisfies ``||grad(f)||_inf <= gtol``.
```

Code says:
```python
tree_mul_term = self.xtol * (tree_l2_norm(params) - self.xtol)
return jnp.all(jnp.array([
tree_inf_norm(state.gradient) > self.gtol,
tree_l2_norm(state.delta) > tree_mul_term
]))
```

Additionally, rather than `all(array(...))` you should use `jnp.bitwise_and(..., ...)` or | & and ~ ops.

## My suggestion
Upon reading up about `madsen-nielsen` stopping condition it seems that there is no single version of it. From my optimisation work I find incorporating both absolute and relative tolerance in parameter changes is quite useful. (Currently it looks like it's only relative)

```python
def leaves_vec(tree_x):
return jnp.concatenate(tree_leaves(tree_map(jnp.ravel, tree_x)))

atol_cond = jnp.all(jnp.abs(leaves_vec(state.delta)) <= self.atol)
rtol_cond = jnp.all(jnp.abs(leaves_vec(state.delta)) <= self.rtol * jnp.abs(tree_vec(params)))
grad_cond = jnp.max(jnp.abs(leaves_vec(state.gradient))) <= self.gtol
done = atol_cond | rtol_cond | grad_cond
return ~done

# defaults
atol = 0. # effectively turned off unless user wants it on, to be backward compatible with current.
rtol = 1e-3
gtol = 1e-3
```

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.