Swapping selected layers with different layers.

[aneesh-aparajit]

Hello!

I am pretty new to Jax/Flax. I was trying to implement to LoRA, although I could modify the state dict, would it be possible to modify the layers to accommodate the additional matrices on run time? Or, is there any other alternate approach to this, so that I could freeze the original weights and just keep the $A$ and $B$ matrices trainable?

I am able to modify the state dict by doing something like this:

import jax
import jax.numpy as jnp

PRNGKey = jnp.ndarray

def get_param(path, param, rank: int, rng: PRNGKey):
    if len(param.shape) == 1:
        return param
    
    a_dim, b_dim = param.shape
    return {
        "a": jax.random.normal(rng, shape=(a_dim, rank)),
        "b": jnp.zeros(shape=(rank, b_dim)),
        "w": param
    }

new_params = jax.tree_util.tree_map_with_path(get_param, params, is_leaf=None)

Is there any maneuver I could do? Or, would I have to make drastic moves like re-defining the architecture?

Any suggestions would help!

Thank you!

[davisyoshida]

I have a JAX implementation of LoRA here: https://github.com/davisyoshida/lorax

You don't need to touch model code to support it.

[aneesh-aparajit]

Thanks a lot for your reply!
The thing is I wanted to learn JAX by implementing something... I did go through your implementation, and spent most of the time with your qax package, which I didn't really understand (me being a beginner). Since, even lorax had a lot of functionalities being derived from qax.

So, my main query is if I were to do it myself, how should I go about implementing it?

Thanks again!

[davisyoshida]

Ah yeah so if you want to really customize JAX you need to write custom interpreters. The JAX tutorials are pretty good (11/10 if we're comparing to most projects, but not 100% of the information you need to be successful).

This is the main one which is relevant: https://jax.readthedocs.io/en/latest/autodidax.html

That being said, if you're looking for something simpler, you can probably write a more restricted version without getting that far down into the weeds. I think targeting only flax Dense layers and making the modification at initialization time should be much simpler. I don't know the best way to do that using FLAX's machinery since I'm a bit more familiar with haiku.

[aneesh-aparajit]

Yeah, the most simplest method would be to just replace the Dense Layers with something like LoraDense or something like that, and then since we can easily modify the parameters we could do that. But, I'll go through the docs as well.
Thanks a lot!