What Are Intrinsic Dimensions? The Secret Behind LoRA



Table of Contents

What Are Intrinsic Dimensions?

Efficiency: Focusing on intrinsic dimensions allows for a significantly reduced parameter count, leading to faster training times and lower computational requirements. This is particularly advantageous for large language models, where parameter spaces are vast.
Fine-tuning: Pre-training large language models implicitly reduces their intrinsic dimension, explaining the success of fine-tuning. These models require less data to adapt to new tasks due to their minimized intrinsic dimension.
LoRA Foundation: Intrinsic dimensions form the basis of Low-Rank Domain Adaptation (LoRA), a technique for efficiently adapting large language models to specific domains. LoRA leverages the concept to achieve rapid development of specialized language models without extensive data.

😎 Intrinsic Dimensions Paper Methodology



Typical Paradigm: If a neural network has $D$ parameters, then we compute the gradient in the entire parameter space $\mathbb{R}^D$.
Intrinsic Paradigm: If a neural network has $D$ parameters, then we choose a random $d$-dimensional subspace ($\mathbb{R}^d$) of the full parameter space and optimise in this subspace. Typically, $d << D$.

👨‍💻 Coding Intrinsic Dimensions

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

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)

We need to flatten the original gradients in order to obtain the shape of the gradients (in this case the grads_treedef).
We take a random normal subspace of the gradients and change its shape to be compatible with the gradients using the tree_unflatten utility function.
In this case, grads are a PyTree object, whereas the intrinsic_space is Traced Jaxpr object. In order to multiply we need to obtain their values which is why jit the element-wise multiplication within tree_map.

📊 The Results



🎬 Conclusion