Random Subspace Training using a custom GradientTransformation

[SauravMaheshkar]

The above mentioned Colab Notebook contains a demo for training a MLP based model for Image Classification on the MNIST dataset.

I wrote a custom GradientTransformation object for training on Random Subspaces of a objective landscapes for a reproducibility study of the paper "Measuring the Intrinsic Dimension of Objective Landscapes" by Chunyuan Li, Heerad Farkhoor, Rosanne Liu, Jason Yosinski.

I wrote it in a way that it can be chained together with other transformations such as

optimizer = combine.chain(
    optax.adam(learning_rate),
    random_subspace_gradients(
        subspace_dim = 16
    )
)

More than happy to make a PR if the devs find it interesting ! Would also appreciate any feedback !

from typing import NamedTuple

from optax._src import base

class RandomSubSpaceState(NamedTuple):
  """State containing PRNGKey for `random_subspace_gradients`."""
  rng_key: chex.PRNGKey

def random_subspace_gradients(
    seed: int,
    subspace_dim: int
) -> base.GradientTransformation:
  """Computes Gradients along some random subspace

  References:
    [Li et al, 2018](https://arxiv.org/abs/1804.08838)

  Args:
    seed (int): initial seed used for the jax.random.PRNGKey
    subspace_dim (int): number of subspace dimensions

  Returns:
    A `GradientTransformation`.
  """
  def init_fn(params: base.Params) -> RandomSubSpaceState:
    del params
    return RandomSubSpaceState(rng_key=jax.random.PRNGKey(seed))

  def update_fn(grads, state, params=None):
    del params
    # split rng
    keys = jax.random.split(state.rng_key, 2)

    # Flatten gradients into 1d array
    grads_flat, grads_treedef = jax.tree_util.tree_flatten(grads)

    # Choose a random intrinsic subspace and reshape as per grads
    intrinsic_space = jax.random.normal(keys[1], (len(grads_flat),))
    intrinsic_space /= jnp.linalg.norm(intrinsic_space)
    intrinsic_space = jax.tree_util.tree_unflatten(grads_treedef, intrinsic_space)

    # project gradients along intrinsic dimension
    projected_grads = jax.tree_util.tree_map(lambda x, y: jax.jit(lambda a, b: a * b)(x, y), grads, intrinsic_space)

    return projected_grads, RandomSubSpaceState(rng_key=keys[0])

  return base.GradientTransformation(init_fn, update_fn)

[fabianp]

Thanks for sharing @SauravMaheshkar !

[SauravMaheshkar]

Thank you @fabianp. Do you think there's interest in adding this to the contrib submodule ?

[fabianp]

Yes, I think it could be a nice addition, specially if it comes with a companion example that reproduces one of the figures in the original paper :-)

[SauravMaheshkar]

Sounds good. I wrote this article (published on Fully Connected) with a Image Classification Example but I'll try and reproduce results from the paper as well. In the middle of exam season atm will keep this discussion thread posted with updates.

[fabianp]

Amazing, and of course no rush!