Integrate Lily

[yibozhong]

Feature request

This request proposes integrating Lily (Low-Rank Interconnected Adaptation across Layers), accepted to ACL 2025 Findings, into the PEFT library.
Paper: https://arxiv.org/pdf/2407.09946
Repo: https://github.com/yibozhong/lily

Motivation

Lily aims to directly make the rank of each individual adapter bigger under the same parameter budget, as it's shown in many papers that higher ranks are beneficial to PEFT performance. This is achieved by breaking the pair-AB-per-layer constraint of LoRA. That is, we do not give each layer a dedicated pair of A and B. Rather, we decouple all the Bs from the layer, and when adapting at each layer, we use a weighted sum of these Bs as the B for this layer. The weight is calculated by a lightweight trainable router, currently data-dependent.

Several points worth noting:

• The method looks somewhat similar to MosLoRA in structure, but it operates at the model level and the aim is to increase the individual rank of each adapter with dynamic adaptation.
• Currently in the paper, we use a data-dependent router, which makes it tricky to merge the weights. I do not observe notable inference latency, possibly due to small model size, but an option for using a non-data-dependent router can be included and enable easy merging the weights.
• The current As are still positioned at a fixed layer (using layer-wise sharing to reduce params). However, it also can be decoupled, simply by providing two routers for weighting As and Bs respectively, rather than one router for B in the current setup. This is a more elegant design and shares the same principle as Lily. After I run quick experiments demonstrating its effectiveness, I can integrate this setup into my current code as Lily v2.

Your contribution

Implement Lily, repo: https://github.com/yibozhong/lily.

[BenjaminBossan]

Thank you for proposing to add Lily to PEFT. Generally, after quickly looking at the paper, I think it could be a nice addition. As you mentioned, the data-dependent routing makes things a bit more complicated. Could you please point me to the part of the code where this is implemented so that I can gain a better understanding?

we use a data-dependent router, which makes it tricky to merge the weights.

Note that the possibility of merging is not a strict requirement for PEFT, so if it's not supported, it's okay.

After I run quick experiments demonstrating its effectiveness, I can integrate this setup into my current code as Lily v2.

Ideally, we would like a new method to be "finalized" when it's being added to PEFT. We had some issues with backwards incompatible changes in the past, which can be especially tricky if they affect the checkpoints. Most of the times, there are workarounds to avoid BC breaking, but they can be quite cumbersome to implement and add tech debt. I'm not sure what your time expectation is, but ideally you can finish the experiments first and then we add the finalized Lily implementation.

[yibozhong]

Thank you for your comments.

Could you please point me to the part of the code where this is implemented so that I can gain a better understanding?

Sure, below is my current impl for the data-dependent routing process. In my repo, refer to this line: https://github.com/yibozhong/lily/blob/main/subject-driven-generation/peft/tuners/lily/layer.py#L146

Note that in the repo, A and B in LoRA are named as lp and hp. I change it back to A and B for better readability.

def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
        self._check_forward_args(x, *args, **kwargs)
        adapter_names = kwargs.pop("adapter_names", None)

        result = self.base_layer(x, *args, **kwargs)
        torch_result_dtype = result.dtype
        lily_A_keys = self.lily_A.keys()
        for active_adapter in self.active_adapters:
            if active_adapter not in lily_A_keys:
                continue
            lily_A = self.lily_A[active_adapter]
            lily_B = self.lily_B[active_adapter]
            B = einops.rearrange(lily_B.weight, "(e i) o -> e i o", e=self.num_B)
            dropout = self.lily_dropout
            scaling = self.lily_alpha
            x = x.to(lily_A.weight.dtype)
            hidden = self.lily_A(x)
            router_logits = self.lily_router(hidden) # [B, N, num_of_experts]
            router_probability = F.softmax(router_logits, dim=-1) # [B, N, ne]
            expert_probabilities = router_probability.mean(dim=(0, 1)) 
            combined_B = torch.einsum("e,eio->io", expert_probabilities, B)
            delta = torch.matmul(hidden, combined_B)
            result = result + (delta * scaling)
            result = result.to(torch_result_dtype)

        return result

Note that the possibility of merging is not a strict requirement for PEFT, so if it's not supported, it's okay.

Thank you for pointing out. Then I can integrate the exact version as descrbied in the paper and repo.

Ideally, we would like a new method to be "finalized" when it's being added to PEFT.

Thanks for pointing this out. Currently, Lily is indeed finalized with the paper being accepted. My mention of another possible design is simply a natural extension of Lily in my mind, which actually differs a lot from Lily and currently lacks experiments. Therefore, I think I can just simply integrate Lily as described in the paper and the repo. As for the alternate design, I think it should be left for the future after I run solid experiments and publish a paper to prove its effectiveness, and maybe then integrate it as a standalone method to avoid any compatibility issues :)

[BenjaminBossan]

Thanks for isolating the core logic. I think it should work with the rest of PEFT, so I'd say you can go forward with your proposal of implementing Lily as currently published.

A small addition: In the snippet, dropout, alpha, and the router are shared for each adapter. I think it would be better to have one per adapter. So:

-           dropout = self.lily_dropout
+           dropout = self.lily_dropout[active_adapter]

etc. WDYT?

[yibozhong]

Thanks for pointing it out. I’ll implement it as you suggested and open a draft PR once it’s done :)

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

[githubnemo]

not stale

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

[yibozhong]

not stale. might delay the integration a bit due to personal matters.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

[githubnemo]

Probably not stale