[RFC] Proposal for complex-valued optimization in Optax

[wdphy16]

As mentioned in #144 , we would like to help enhance the support for complex numbers in the JAX community. The detailed proposal document is here: https://gist.github.com/wdphy16/118aef6fb5f82c49790d7678cf87da29

In brief, we need to properly implement the norm of complex variables in the optimizers. We need to decide whether to implement the complex norm, the split real norm, or both of them.

Although there is also a comment zone below the gist, I would like to keep the discussion in this issue thread. Feel free to leave your comments!

[rosshemsley]

Hello @wdphy16, thanks for your comment, this initiative is interesting and welcome!

One consideration is that optax is heavily used by models that are 'internal', and so any changes that we pull upstream have to be carefully managed to avoid performance penalties, semantics changes, and drastic API changes that could cause problems to existing optax users.

We do believe that these ideas can be integrated whilst managing this complexity though, and think this could be an important set of features for optax. Perhaps a good way to start could be to sketch up a proposal document, suggesting a concrete approach - effectively a more detailed version of your comment, where we can add suggestions. That way we can help to avoid some of the pitfalls given above up front :)

[wdphy16]

Hi @rosshemsley , thanks for your feedback! Should I write the proposal document here, or in a PR or Discussions?

[rosshemsley]

Hey @wdphy16: perhaps an issue thread, or a github gist?

Any approach that is publicly readable and comment-able would be fine,

One example of a proposal with public comments is here: golang/go#43651 (consider this an example, though probably this is far more detailed than would be needed in this case).

[wdphy16]

Hi @rosshemsley , sorry for the late reply. I've written the proposal document and edited the first comment of this issue.

[mtthss]

Thanks @wdphy16 for sharing such a thoughtful and detailed proposal.

I think we should certainly implement the split real norm approach described in your doc.
That is an enhancement of the current capabilities of optax with literally no downsides,
as those who are not interested in complex number will just continue not to use the split_complex wrapper.

Could you put together a PR for this?

About the other proposal, I think it would be nice to support that as well,
but we need to think carefully about performance and readability,
for instance it would be good to verify for sure what XLA is or isn't able to do

E.g. the point from your doc:

When g is real, it should be possible for jax.jit to eliminate the dispatch overhead of conj and real, and optimize (g * g) ** (order / 2) into g ** order (TODO: I'm not sure if currently jax.jit is smart enough to do this)

[wdphy16]

It turns out that currently jax.jit cannot optimize the composition of a multiplication and a power. I did some quick experiments with this debugging function make_hlo:

In [1]: def norm(g, order):
   ...:     return (g.conj() * g).real ** (order / 2)
   ...:

In [2]: g = jnp.array([1., 2., 3.])

In [3]: print(make_hlo(norm)(g, 2))
HloModule xla_computation_norm.11

ENTRY xla_computation_norm.11 {
  constant.3 = pred[] constant(false)
  parameter.1 = f32[3]{0} parameter(0)
  multiply.4 = f32[3]{0} multiply(parameter.1, parameter.1)
  parameter.2 = s32[] parameter(1)
  convert.5 = f32[] convert(parameter.2)
  constant.6 = f32[] constant(2)
  divide.7 = f32[] divide(convert.5, constant.6)
  broadcast.8 = f32[3]{0} broadcast(divide.7), dimensions={}
  power.9 = f32[3]{0} power(multiply.4, broadcast.8)
  ROOT tuple.10 = (f32[3]{0}) tuple(power.9)
}

In [4]: print(make_hlo(norm, optimize=True)(g, 2))
HloModule xla_computation_norm__1.11

fused_computation {
  param_1.4 = f32[3]{0} parameter(1)
  multiply.2 = f32[3]{0} multiply(param_1.4, param_1.4)
  param_0.1 = s32[] parameter(0)
  convert.0 = f32[] convert(param_0.1)
  constant.1 = f32[] constant(0.5)
  multiply.1 = f32[] multiply(convert.0, constant.1)
  broadcast.0 = f32[3]{0} broadcast(multiply.1), dimensions={}
  ROOT power.0 = f32[3]{0} power(multiply.2, broadcast.0)
}

ENTRY xla_computation_norm__1.11 {
  parameter.2 = s32[] parameter(1)
  parameter.1 = f32[3]{0} parameter(0)
  fusion = f32[3]{0} fusion(parameter.2, parameter.1), kind=kLoop, calls=fused_computation
  ROOT tuple.10 = (f32[3]{0}) tuple(fusion)
}

We can see that in the optimized HLO conj and real are indeed eliminated for the real g, but there are still a multiplication and a power. Actually such optimization does not exactly preserve the floating point arithmetic result, so it should be enabled only with some dangerous flag like 'fast math', and I don't know if JAX developers will have such plan in the future.

Anyway, we can still use an if-else statement to eliminate that overhead when g is real. I've updated the proposal.

[PhilipVinc]

@wdphy16 can you update the PR to use the if-else statement and fix ADAM? I think this should be good for @mtthss