[Feature Request] Normalized gradient descent

[smorad]

Optax has various clipping operators, but as far as I can tell, it cannot scale by gradient norm. Adding these capabilities such that they could be chained would allow us to use normalized gradient descent methods (e.g. normalized Adam, etc).

A simple implementation might look like

def scale_by_norm(scale: float=1.0, eps: float=1e-6):
  def init_fn(params):
    del params
    return optax._src.base.OptState

  def update_fn(updates, state, params=None):
    del params
    g_norm = jnp.maximum(optax.global_norm(gradient) + eps, scale)
    def scale_fn(t):
       return t / g_norm

    updates = jax.tree_util.tree_map(scale_fn, updates)
    return updates, state

  return optax.GradientTransformation(init_fn, update_fn)

[mtthss]

Do you have a reference to this specific way of normalising?

[smorad]

This textbook describes it fairly well. My example might be a little fancy, but you could replace the maximum with

g_norm = (optax.global_norm(gradient) + eps) / scale

In this case, scale would refer to alpha in Eq 6.

[mtthss]

Sounds like it could be a good addition. Do you want to put together a PR?

[SauravMaheshkar]

Seems like a simple extension of

optax/optax/contrib/sam.py

Lines 63 to 80 in 841be5a

	def normalize() -> base.GradientTransformation:
	"""Normalizes the gradient.
	Returns:
	An (init_fn, update_fn) tuple.
	"""
	def init_fn(params):
	del params
	return NormalizeState()
	def update_fn(updates, state, params=None):
	del params
	g_norm = utils.global_norm(updates)
	updates = jax.tree_map(lambda g: g / g_norm, updates)
	return updates, state
	return base.GradientTransformation(init_fn, update_fn)

@mtthss can I take this up ?

[smorad]

I think this might actually be implemented in clip_by_global_norm. IIRC the code there actually scales the gradient rather than clips it. Might be worth double checking before starting.

[vroulet]

clip_by_global_norm clips but do not necessarily normalize (if the updates are less than clip norm, then they are just returned as is). In other words clip projects on a ball and @smorad you want to project on a sphere.
I think @SauravMaheshkar pointed out a good starting point.