Fine-Tune model extraction

[krahnikblis]

loving this work!

so, wondering, hoping folks smarter than i can weigh in: the concept of LoRA is that it's wrapping/replacing some layers in the model, and basically applying some extra weights of its own to the existing params/state dict, right? W + delta-W from the front page... so basically, existing model + LoRA method with weights = new model of the same effective structure?

so my question is, and i know i'm oversimplifying, can we not "extract" the differences between fine-tuned models and the starting points, save that difference as a LoRA? the "training" process would not need images or the text/vision processing, just 2 sets of checkpoints and some loops to get the weights right, such that original model + LoRA = fine-tuned model. i.e., plug in SD1.4 as original, @nitrosocke's (flippin awesome btw) nitro-diffusion as the target, run the maths, and then you have a LoRA compression representing all the learnings of arcane, mo-di, archer, but in a few MBs. do the same for your other favorite fine-tunes, mix/match from there.

if this works it would also save the facehuggers all those PBs of storage being taken up by everyone and their dog's fine-tune models (so many, just on butterflies alone)...

[hdon96]

@brian6091 showed in his discussion post here that lora trains a subset of the dreambooth parameters whose size is based on lora rank, so I'm not sure the same model could be easily transformed from dreambooth into a lora. Reproducing a dreambooth model by training lora with the same dataset and similar parameters seems effective so far.

[cloneofsimo]

Hmm. This might actually work. We might low-rank approximate the "difference" of a model and start from there. Not sure how this would work on "heavily fine-tuned models".
BTW, we can, in a numerically stable way, explicitly find such low-rank matrix via SVD.
https://en.wikipedia.org/wiki/Low-rank_approximation#Basic_low-rank_approximation_problem

[brian6091]

Cool, I understand the idea. Might generate some interesting effects, I think particularly for styles.

On a side note, interesting to know what model weights are changing most for certain concepts, e.g., people versus styles. I wonder if someone has visualized where say the top million weights with the highest weight change are distributed (or look at big gradients early on in training).

[krahnikblis]

this is sorta the stuff i do to grok things. like, i am currently trying to look at diff of unet pre/post lora using this. but i'm not seeing a difference here, i might've done something wrong. and with textual inversion, i use pyplot to output the embedding vector before, diff, and after, at each epoch, so i can see the distribution of change over the vector's values. it would be harder to do this with multi-dimensional weights, but i was also going to tinker with interpolating (with some sort of spline method, not linear) a new set of weights, given the diffs among sd1.x and the various fine-tunes.

[brian6091]

Nice, thanks for the torchview pointer, looks like fun. Your approach sounds smarter than what I'm doing, which is just dumping weights out to the console...

[brian6091]

this is sorta the stuff i do to grok things. like, i am currently trying to look at diff of unet pre/post lora

@krahnikblis
thought you might appreciate this figure from this paper that @cloneofsimo pointed out in another thread:

https://drive.google.com/file/d/1yVwVh1TBgpjVaoDdMZZA-XCPShoI_gQu/view?usp=share_link

[krahnikblis]

noice! i probably should read the paper to find out, but when we talk about 'delta' here, sometimes the delta is just down to noise, so it would be interesting to zoom in on the important parts of the deltas. basically, these models to me feel like big Plinko boards - the whole board doesn't matter, what matters is the part of the board that forms the path your ball takes from top to bottom. this starts to get into some of my other musings about on-the-fly language-aware weight pruning

but, what i was trying to articulate re diff was, does the Lora Injected Layer show up as a difference in the model's structure? i think the default value for the depth arg in the torchview library is not enough to find out, and i just haven't done it to a deeper level yet. basically i wanted to understand how the layer wraps/replaces/augments the existing model's modules, so that i can understand how to approach the JAX version. i think i get it now, so the next part is how to edit pytrees to do the equivalent, and how to make just the LoRAs the trainable state... today i've integrated the dreambooth and textual inversion Flax scripts, and am now working on replicating the LoRA parts.

[krahnikblis]

awl hell nawl. here you go:

# plot tokens function
def plot_token_embed_diff(orig, new, dpi=96):
    plt.rcParams["figure.dpi"] = dpi
    plt.rcParams["figure.figsize"] = orig.shape[-1] / dpi, 0.5
    plt.style.use('dark_background') 
    # ref https://stackoverflow.com/questions/45841786/creating-a-1d-heat-map-from-a-line-graph
    # cmaps https://matplotlib.org/stable/gallery/color/colormap_reference.html
    fig, (ax1,ax2,ax3) = plt.subplots(nrows=3, sharex=True)
    xmax = orig.shape[0]
    # orig on top:
    ax1.imshow(jnp.expand_dims(orig,0), cmap='gray', interpolation='nearest', aspect="auto", extent=[0,xmax,0,6])
    ax1.set_xticks([])
    ax1.set_yticks([])
    # diff in middle:
    range = np.maximum(np.abs(np.min(new - orig)),np.max(new - orig))
    ax2.imshow(jnp.expand_dims(new - orig,0), cmap='jet', interpolation='nearest', aspect="auto", extent=[0,xmax,0,6],vmin=-range,vmax=range)
    ax2.set_xticks([])
    ax2.set_yticks([])
    # new on bottom:
    ax3.imshow(jnp.expand_dims(new,0), cmap='gray', interpolation='nearest', aspect="auto", extent=[0,xmax,0,6])
    ax3.set_xticks([])
    ax3.set_yticks([])
    plt.show()

i just pop the (768,)-shaped vectors in there after each epoch. the output looks like a pair of ... hmm maybe i can just paste a sample here...

sweet it worked! so yeah, the middle one is a heatmap of the delta over the epoch. instead of a bar chart for a single-dimension vector, i guess you could adapt this but plot 2D versions, kinda like images, although it might get hard to view that way too...

[brian6091]

Thank you!

[cloneofsimo]

In many cases, low-rank is too low to approximate stuff, but when lucky, it fits some parameters very well. I think we should regard this two-fold. First pruning, second knowledge-distillation.

[cloneofsimo]

To be honest, I'm surprised that this merely worked. What?
here is text prompt
"redshift style, cute pokemon cat"
with redshift diffusion

Here is base sdv.1.5

Now (surprise surprise), here is

$$ Z = W_{redshift} - W_{base} $$

$$ U, V = \arg \min _{U, V} || U V - Z ||_F , rank(U) = rank(V) = r $$

$$ W_{redshift approx} \leftarrow W_{base}+ 3.5 UV $$

What?? This works surprisingly well!! haha

[brian6091]

Well, color me surprised. Crazy.

[cloneofsimo]

Last one is very cute in my opinion

[krahnikblis]

SWEET! it looks like it works well enough! this could be very valuable for working toward small-footprint models and customizations

anyway, to my dumb little marketing brain, it sounds like this diff-svd-lra stuff could be getting into the realm of something like this paper. i'd like to catch up and try these things myself, but i'm currently trying to be disciplined and stay focused on translating the torch stuff to JAX and get it running optimally on TPUs. so again, thanks for testing this idea so soon! glad it's fruitful!

[krahnikblis]

@cloneofsimo could you provide the code snippet you used to get the model weights diff and apply SVD to get the lora up/down weights?

[krahnikblis]

ok the reason i ask for the code snippet for SVD test is: i'm a career marketer and it's been 20 years since college-level maths, so this is slow learning for something that's peripheral to my objective (which is performant custom art tools)... and the other reason is, it seems like this SVD method is required beyond approximating existing fine-tunes, because...
the approach i'm taking is, instead of hacking the model itself (i.e. replacing the modules with the LoraInjectedLinear), to apply the "lora model" directly to the parameters, such that the result of it is the adjusted/finetuned params. and the reason for doing this is also 2-fold:

1. jax models treat parameters separately from the model, so it would need double the transforms/codework to adjust params and module separately. (i work with jax for multiple reasons: performance, skilling up for relevant numpy projects i'd personally like to accelerate, and personal opinions/frustrations about pytorch.)
2. it's my hypothesis that a standalone lora-on-param model would be more versatile regardless of jax/torch (e.g. my original question in this thread), and more robust to changes over time (e.g. it's VERY often for the facehuggers to alter transformers or diffusers, which might break a LoRA module's expectations - but applying LoRA to the weights means that wouldn't happen)

EDIT: i had a bunch of confusing questions and a code sample here that went the wrong direction, which would not contribute value to future readers. so i'm replacing it with net learnings from my further tinkering, and cloneofsimo's replies below:

1. SVD method produces U, S, Vh - when these are dot-producted back together, you get your original array.
2. you can reduce the rank of U and Vh (alternately, you can test this by zeroing-out S for indexes grater than desired rank) to get the low-rank approximation. still, when you dot-product the lower-rank arrays back together, you get something the same shape as original array, with the most "important" attributes still in it, i.e., the approximation
3. when low-ranked, U * S and Vh, or some variation of maths thereof (e.g. i'm using U * sqrt(S) and diag(sqrt(S)) @ Vh), are the "lora_up" and "lora_down" params.
4. in other words, you can dot-product the lora up/down params to get the approximate original array. i got really lost in the implementation from this library's "injection", which basically does the dot product of the difference but wraps it in a combined method that ALSO performs the main model's attention layers' linear transforms. once i realized points 1-3 here, i understood that we don't need to do it that way, i.e. my hypothesis in original point 2 above holds true - we can have a standalone low-rank model that "wraps" the main model's params. this works great in JAX, but i don't know how it would translate back to torch.
5. answer to this entire thread question is "YES", we can "extract" the meaningful part of a fine-tuned model by taking the low-rank SVD of the difference of just the attention layers. crazy that it works, and works well.

[cloneofsimo]

Here is my code to make it work:

def extract_linear_weights(model, target_replace_module):
    lins = []
    for _module in model.modules():
        if _module.__class__.__name__ in target_replace_module:

            for name, _child_module in _module.named_modules():
                if _child_module.__class__.__name__ == "Linear":
                    lins.append(_child_module.weight)

    return lins


ori_unet = extract_linear_weights(pipe_base.unet, ["CrossAttention", "Attention", "GEGLU"])
ori_clip = extract_linear_weights(pipe_base.text_encoder, ["CLIPAttention"])

tuned_unet = extract_linear_weights(pipe_tuned.unet, ["CrossAttention", "Attention", "GEGLU"])
tuned_clip = extract_linear_weights(pipe_tuned.text_encoder, ["CLIPAttention"])

diffs_unet = []
diffs_clip = []

for ori, tuned in zip(ori_unet, tuned_unet):
    diffs_unet.append(tuned - ori)

for ori, tuned in zip(ori_clip, tuned_clip):
    diffs_clip.append(tuned - ori)
    
 uds_unet = []
uds_clip = []
rank = 4
with torch.no_grad():
    for mat in diffs_unet:
        mat = mat.float()
  
        assert hi_val > low_val
        
        
        U, S, Vh = torch.linalg.svd(mat)

        U = U[:, :rank]
        S = S[:rank]
        U = U @ torch.diag(S)
        
        Vh = Vh[:rank, :]

        dist = torch.cat([U.flatten(), Vh.flatten()])
        hi_val = torch.quantile(dist, 0.999)
        low_val = -hi_val
        
        
        U = U.clamp(low_val, hi_val)
        Vh = Vh.clamp(low_val, hi_val)
        
        uds_unet.append(U)
        uds_unet.append(Vh)
    
    for mat in diffs_clip:
        mat = mat.float()
        
        U, S, Vh = torch.linalg.svd(mat)

        U = U[:, :rank]
        S = S[:rank]
        U = U @ torch.diag(S)
        
        Vh = Vh[:rank, :]

        dist = torch.cat([U.flatten(), Vh.flatten()])
        hi_val = torch.quantile(dist, 0.999)
        low_val = -hi_val
        
        U = U.clamp(low_val, hi_val)
        Vh = Vh.clamp(low_val, hi_val)

        
        uds_clip.append(U)
        uds_clip.append(Vh)
        
  from lora_diffusion import (
    monkeypatch_lora,
    monkeypatch_add_lora,
    monkeypatch_replace_lora,
)

monkeypatch_lora(pipe_base.unet, uds_unet, ["CrossAttention", "Attention", "GEGLU"])
monkeypatch_lora(pipe_base.text_encoder, uds_clip, ["CLIPAttention"])

Ok so obviously very hacky code and might not work with recent update. I'll definitely share the cleaned up version with #88 release.

[cloneofsimo]

And sorry for not being helpful here : I'm not really familiar with jax, so I wouldn't be much help... Also agree 100% on the versatile huggingface nature. can (and will) break LoRA in the future. In that case though, We would only need to fix find_module function, which would be pretty easy fix. LoRA codebase is also kinda tiny, so I'm ok with rewriting the code haha

[krahnikblis]

thanks for the code! i also got the thing above working (using SVD and LoRA params in a model that wraps the target params, not injects to the target model). whereas before i was getting weird noise and gray outputs in images, now i'm getting art! but, for the same test prompts, i get drastically different styles/images for kiriko LoRAs vs what y'all have been posting, so i must have missed something somewhere... BUT! success in workflow and figuring out the linalg stuff!

JAX (which is basically numpy, scipy, and autograd and neural network utilities) on TPUs is nice, you can put 8 prompts at once, or 8 different guidance scales, or whatever, and get 8 images at once with render time of ~30sec. (8 is number of TPU cores in Colab). i did get some issues with JAX's implementation of SVD - something about incompatible data types, but it only happened on the unet, so i suspect it's something with where the arrays are located - i suspect the transformer vs diffuser processes for creating params are different. anyway, workaround was just to cast the dtypes before and after "lorafying" the unet params.

[krahnikblis]

@cloneofsimo regarding your SVD code above (and THANKS btw, that set me on a path to better learn what SVD and LoRA are mathematically), i'm poking around with the outputs, and noticed a few different things:

1. in the existing checkpoint lora_kiriko.pt, all the lora_up weights for all to_out modules are 0! not the down-weights, so i don't think this is connected to the to_out 0/1 issue from couple weeks ago... does this mean the to_out modules can be completely ignored, with no loras calculated for them? i assume that with up weights == 0, the entire lora_up(lora_down(to_out)) result would be 0...
2. creating up and down weights from SVD U*S and Vh... i was wondering why you chose U to multiply by S, instead of Vh, or instead of taking root of S and multiplying both by that. real-life data would point to using down = S * Vh: in all the lora weights for kiriko, the down weights are the ones with larger (magnitude) sums and means, while in my experiment with SVD, the U values' sums and means are already generally larger than the Vh. i.e., should the choice of where to multiply S be reversed? i know it all gets multiplied anyway, so maybe doesn't matter at all...
3. what mathematical wizardry causes LoraInjectedLinear to 'automatically' find the equivalent to the highest singular value vectors? i don't see ranking by importance or anything, so is it that the training process itself is just filling the weights with the best possible via gradient descent? i now wonder which result would be 'better' for the purpose - a true SVD low-rank approx of trained - original weights, or this trainable up-down method? because i'm scratching my head over how to implement trainable up/down in Flax that is a different call vs train process, i think i'll investigate this in the normal training process, but unfreeze only the layers that LoRA would be applied to, then use this SVD to extract the ups/downs...
4. given that U * S * Vh = original array, this means that up_weight * down_weight = approx trained diff right? i.e., all my shape-fumbling code above is completely bonkers and betrays my dumb little marketing brain... just need to add up*down to the base weights... right?
5. connecting questions 3 and 4: could taking SVD low ranks at each train step be an optimal method? i.e. grads are collected for whole layers, but instead of updating weights with them directly, calculate the loras of the grads. could even progressively lower the rank over training epochs, and dynamically choose the best rank for each layer, i think akin to what @brian6091 was asking earlier...

here's the code that shows the 0-valued up-weights for to_out layers:

import torch
import numpy as np
from diffusers import UNet2DConditionModel
sd_unet = UNet2DConditionModel.from_pretrained(sd2_model_name_or_path, subfolder="unet", torch_dtype=torch.float16) 
!wget https://github.com/cloneofsimo/lora/raw/master/contents/lora_kiriko.pt
kiriko = torch.load("lora_kiriko.pt", map_location=torch.device('cpu'))
@torch.no_grad()
def extract_linear_weights(model, target_replace_module):
    lins = []
    i = 0
    for _module in model.modules():
        if _module.__class__.__name__ in target_replace_module:
            for name, _child_module in _module.named_modules():
                if _child_module.__class__.__name__ == "Linear":
                    lins.append(_child_module.weight)
                    print(_module.__class__.__name__,name, _child_module.weight.shape)
                    up = kiriko[2*i].numpy()
                    dn = kiriko[2*i + 1].numpy()
                    print("up", up.sum(axis=0), up.mean(axis=0), "down", dn.sum(axis=1), dn.mean(axis=1))
                    i += 1
    return lins
# this loads the list and prints the stats/names:
sd2_unet = extract_linear_weights(sd_unet, ["CrossAttention", "Attention"]) # , "GEGLU" adds some torch.Size([2560, 320]), and isn't present in kiriko

[cloneofsimo]

Whoa thats a lot.

1. Yes, this was handled in pr about two weeks ago.
2. Yes it doesn't matter which one is being multiplied
3. It doesn't guarantee to align with spectrums. To me at least, it just empirically works and that seems to be the case.
4. Yes.
5. I'm not sure. What you are referring to in general connects to much bigger research topic : manifold optimizations, which in terms of local optimization is about taking optimal gradient in general spaces.

[cloneofsimo]

Just merged a PR #98 that allow you to do this with CLIs!

[bmaltais]

OK... this sound interesting. Is this a method to extract a LoRA network from a finetuned model? For example, one create a dreambooth model from sd1.5... so this method would extract the delta between the base model and the finetuned model and save it as a LoRA? Am I understanding this correctly?

If so... any example of the cli command to do this?

[krahnikblis]

yeah, it works exactly like that, I've already extracted ~10 models. fun part will be to mix a bunch of them together! clonofsimo said the last release includes this but i don't know which cli file it is... i built it separately using Flax in a Colab notebook, focused on TPUs... once i get back into it and clean up my crap i could post my notebook, if anyone's interested. it's not directly backward compatible to pytroch though, since some tensors are transposed and i built it to save as msgpack

[cloneofsimo]

Yes you are understanding this correctly. I will add more info on the readme. Meanwhile, you can use lora_distill command on dev branch to get it work

[bmaltais]

Trying to convert a model to lora... but somehow I can't figure out what the lora_distill is looking fot:

lora_distill "D:\models\bernard\astria-sks_man-sd15.ckpt" runwayml/stable-diffusion-v1-5

Error: OSError: It looks like the config file at 'D:\models\bernard\astria-sks_man-sd15.ckpt' is not a valid JSON file.

How can I extract the LoRA from a CKPT model?

Ohhhh... does it only work with diffuser models? So I need to convert my CKPT to a diffusers 1st

[cloneofsimo]

Yes it only works eith diffuser models for now.

[brian6091]

The CLI file that does the difference decomposition is cli_svd.py on the develop branch. I'm not sure if there's anything for mixing multiple decompisitions together.

[krahnikblis]

i'll add my methods for this once i've tested and refined. but what works for me in the CLIP embeddings (i.e. TI) is to sum them up and take the norm of the sum. however, those embeds expect to be in the 0-1 variance range, but these attention layer diffs don't, so i've got a "norm_to" function that fixes a new normed vector's norm range to some reference vector's range, which will be necessary here. the other thing i need to test is, do we take the sum of the loras themselves? or the sum of their dot products? i'm currently focused on the training process in JAX, so this will be something i get to next...

[wsxiaoys]

Sharing some svd distilled result of prompthero/openjourney

Lora safetensors are created with:
lora_distill prompthero/openjourney runwayml/stable-diffusion-v1-5 --rank 8

I used scale 1.0 for both unet / text encoder in distilled model inference.

Openjourney

SVD Distilled Openjourney

Openjourney

SVD Distilled Openjourney

Openjourney

SVD Distilled Openjourney

[cloneofsimo]

Thats cool haha

[krahnikblis]

nice! i also found success here, and built a combinator with lots of different mathsy methods to test combo methods. example, these are Arcane (from Nitro diffusion) and Tron Legacy extractions combined...

[mekaneeky]

I am utilizing an extension of this SVD-distillation approach in a paper I am soon releasing. Do we know who pioneered SVD distillation and how can I cite this ? Do I just point to this thread?

[krahnikblis]

@cloneofsimo demonstrated the SVD extraction as a response to my OP question, and i think i recall the paper also describing, or referencing the Microsoft paper (which i think was about LLMs), using SVD to essentially compress weights. so i don't know if "pioneer" is anyone here, but either/both of those papers make the point of using LRA on attention layer weights to make fine-tuning orders of magnitude more portable...

i've been doing a lot of tinkering with how to combine and "re-inflate" the SVD/LoRA extractions i've taken from multiple models. the weights of models seem to be finicky, and sometimes i can get really lucky with certain combinations of both the models i select and the method i use to combine them (e.g., i've got 6 models i like that i can simply sum the loras and then use dot-product to inflate to original parameter size, and then add those to the SD1.5 base; but, if i try the same lora-sum method to 4 other models i like, everything is completely fried and i have better results with more sophisticated maths...). i'm currently trying to run some controlled tests to get to a recipe that always works, and a way to predict what would work best in what scenario...

all this is to say, i'd like to read your paper and if you wanna geek out on this stuff and teach me a thing or 2 about matrix math and machine learning, i'm game...

[cloneofsimo]

Ive worked on this for quite a while after initial findings, and developed a method coined ere. Quite interestingly, low rankedness is rather unique phonomena of weight residuals which is what makes this possible. Perhaps you can cite this work. https://arxiv.org/abs/2305.18425

[cloneofsimo]

I too havent found a single paper that tackled this. It is completely possible though, because thats why PEFT was developed in the first place, and fine tuning + compressing residual isnt a natural choice. We tackle this problem in our paper as well.

[mekaneeky]

Sorry, just saw your posts now. I developed a small extension that I haven't released yet (writing, responding to reviewers, etc) should stop procrastinating and submit a version to arxiv in a day or 2, yet is very similar to this work SVDdiff That I recently discovered and I think is concurrent to mine, and slightly more clever I have to admit (they parameterize the singular values only, leading to very small sizes).

I have been thinking about the fact that we store neural network weights as 2d matrices. Is that only due to convenience? or can we come up with more efficient structures for compressing neural network weights (that can be decompressed to 2d matrix form, or used as is), ideas are still not very mature. But playing around.