Bug init original SAE hidden states with hypernetwork

[thangld201]

Hi @zedian, in algorithm 1 of the paper the SAE representations (from RoBERTa model, without hypernetwork) is used in training. However, in alignment.py, when initializing the RobertaLoraWrapper model, it automatically injects (in-place) between the query and value projection layers with low-rank adapters from hypernetworks. This was done before SAE representations are calculated, which doesn't make sense. Could you re-check here ?

In alignment.py:
......................................................................
......................................................................
    if model_args.apply_hyperlora:
        # init model, hypernetwork already here
        model = RobertaLoraWrapper(model, data_args.dialects, model_args.lora_rank, model_args.hidden_adapter_dim, model_args.load_hypernet_weights, train=True) 
......................................................................
......................................................................
    model = model
    dataset = dialect_datasets.map(
        tokenize_text,
        batched=True,
        load_from_cache_file=not data_args.overwrite_cache,
        desc="Running tokenizer on dataset",
    ).map(
        embed_text,
        batched=True,
        load_from_cache_file=not data_args.overwrite_cache,
        desc="Original Embeddings for Untransformed Input",
    ) # calculate original embeddings, but hypernetwork already injected