ICML 2024 DoRA：权重分解低秩适应详解

• Adapter-based: 在冻结的预训练模型中插入额外可训练模块
• Prompt-based: 在输入中添加可训练软提示（soft tokens），仅微调提示向量
• LoRA 及其变体: 用低秩矩阵近似权重更新，推理前可与预训练权重合并

3 LoRA 与 FT 的模式分析（Pattern Analysis of LoRA and FT）

名词理解

• 全量微调（FT, Full Fine-Tuning）
  • 为给训练模型'全面补课'，让已经学过通用知识的模型，针对具体任务打磨每一个'知识点'
  • 特点：模型的所有参数都会改，会根据具体任务调整每个参数的权重
  • 缺点：时间和金钱成本太高，易产生'灾难性遗忘'
• 低秩适应（LoRA, Low-Rank Adaptation）
  • 特点：不更新所有参数，只针对模型适配下游任务需要调整的关键参数进行修改
  • 核心逻辑：用两个小矩阵的乘积代替大矩阵
  • 缺点：和全量微调比存在精度差距

LoRA 数学原理与核心公式

LoRA（Low-Rank Adaptation）通过低秩矩阵分解近似权重更新，在不增加推理延迟的前提下实现参数高效微调，核心公式如下：

1. 权重更新总公式 预训练权重矩阵 $W_0 \in \mathbb{R}^{d \times k}$（$d$ 为输出维度，$k$ 为输入维度），微调后的权重 $W'$ 由'预训练权重 + 低秩更新量'两部分构成： $$W' = W_0 + \Delta W = W_0 + \underline{BA}$$ $W_0$ 在微调过程中保持不变，$\Delta W = BA$ 为低秩更新量（$\underline{BA}$ 表示可训练参数），$B \in \mathbb{R}^{d \times r}$、$A \in \mathbb{R}^{r \times k}$ 为低秩矩阵，且 $r \ll \min(d,k)$（通常 $r = 4, 8, 16$） B 初始值为零，BA 初始值为零。

2. 初始化策略 为确保训练起始时无更新干扰，矩阵 A 采用 Kaiming 均匀分布初始化，矩阵 B 初始化为全零矩阵： $$A \sim \text{KaimingUniform}(\cdot), \quad B = 0 \implies \Delta W|_{t=0} = 0$$

3. 表达能力与秩的关系 当 LoRA 的秩 $r$ 趋近于预训练权重的固有秩 $r^$ 时，其表达能力可接近全量微调（FT）： $$\lim_{r \to r^} \mathcal{E}(W_0 + BA) = \mathcal{E}(W_{FT})$$ 其中，$\mathcal{E}(\cdot)$ 表示模型在下游任务上的损失函数。

3.2 权重分解分析

LoRA 可被视为全微调的一种通用近似，通过逐渐增加 LoRA 的秩，使其与预训练权重的值保持一致，能够达到与全微调相近的表达能力。

先前研究：将两种方法的精度差一归因于可训练参数数量有限，未进行深入分析。

研究人员得到的启发：权重归一化——将权重矩阵重新参数化重构为幅度与方向两个独立的部分——来研究 LoRA 和 FT 学习模式的差异。

分析方法：考察 LoRA 和 FT 权重相对于预训练权重在幅度和方向上的更新，以揭示两者学习行为的根本差异。 $W \in \mathbb{R}^{d \times k}$ 的权重分解可表述为：

$$W = m \cdot \frac{V}{|V|_c} = |W|_c \cdot \frac{W}{|W|_c}$$

其中，$m \in \mathbb{R}^{1 \times k}$ 是幅度（大小）向量，$V \in \mathbb{R}^{d \times k}$ 是方向矩阵，$|\cdot|_c$ 是矩阵各列的向量范数。这种分解确保 $\frac{V}{|V|_c}$ 的每一列都是单位向量（仅反映方向，与大小无关），而 $m$ 中的相应标量表示该列的幅度大小。$|W|_c$ 表示矩阵 W 按列的逐向量范数（即对每一列计算向量范数，反映该列的'强弱程度'）。

通俗理解 简单来说，这个公式就是把权重矩阵分解成'强弱程度'和'作用方向'这两个独立的部分，方便我们分别研究它们的变化规律。

• 方向矩阵（V）：可以理解成'调整的方向'。比如你要把工具往东边挪、往西边调，这个'东、西'就是方向。这里通过把每一列都变成'单位向量'，就像把方向统一规范成'标准朝向'，方便对比。
• 幅度向量（m）：可以理解成'调整的力度'。方向确定了，是使劲调（幅度大）还是轻轻调（幅度小），就由 $\frac{V}{|V|_c}$ 里的标量来决定。
• 而 LoRA（低秩适应）和 FT（全微调）的区别，就体现在**怎么调整这些'幅度'和'方向'**上：
  • LoRA 是'针对性微调'：只调整部分低秩的'幅度 - 方向'组合，就像给工具的某些关键部位小幅度、精准调整。
  • FT 是'全局大调整'：把所有权重的'幅度'和'方向'都重新调一遍，相当于把工具从头到脚彻底改造。

LoRA 精度有时不如 FT——它在'精细调整'上能力不足，比如想大幅调整方向的时候，没法同时灵活调整大小。 既然知道了问题，我们就提出了 DoRA 这个方法：先把预训练权重拆成'大小'和'方向'两部分；对'方向'部分再用 LoRA 来微调（因为方向部分参数多，LoRA 能保证高效）；同时让'大小'部分也能微调。

$W$ 和 $W_{FT}$ 之间的幅度和方向变化可定义如下： $$\Delta M_{FT} = \frac{\sum_{n=1}^{k} |m_{FT}^{n} - m_0^{n}|}{k}$$ 可以理解为'平均差异有多大'。 $$\Delta D_{FT} = \frac{\sum_{n=1}^{k} (1 - \cos(V_{FT}^{n}, W_0^{n}))}{k}$$ 可以理解为'方向差异有多大'。

• 研究人员用这两个公式，分别从'幅度'和'方向'两个角度，对比了不同方法（Sung 的方法和 LoRA）在模型权重上的变化，以此来分析哪种方法的训练效果更优。
• 得出结论：LoRA 缺乏进行更细微调整的细致能力；不擅长在进行较大幅度改变的同时执行细微的方向变化，同时学习幅度与方向是一个挑战——引出下文提出 LoRA 的变体——DoRA！

4 方法（Method）

4.1 权重矩阵低秩适应 权重分解低秩适应（DoRA）将预训练权重分解为幅度分量和方向分量，因方向分量参数较多进一步使用 LoRA 进行分解。（LoRA 仅专注于方向适应，简化）

DoRA 与权重归一化的不同点是：权重归一化从头开始训练两个组件，而 DoRA 开始于预训练权重，避免了初始化敏感问题。

DoRA 的公式： $$W' = m \cdot \frac{V + \Delta V}{|V + \Delta V|_F} = m \cdot \frac{W_0 + B \Delta A}{|W_0 + B \Delta A|_F}$$

DoRA 解决了模型初始化的问题。它先把预训练的权重 $W_0$ 初始化为两个低秩矩阵 B 和 A 的乘积，并将 $m = |W_0|$ 作为可训练的标量参数，通过类似 LoRA 的方式更新方向。 $\Delta V$ 是由矩阵 B 和 A 的增量相乘得到的方向更新。简单来说，就是先拆分预训练权重，再通过可训练参数调整，最后规范权重得到新结果。

这张图展示了三种模型微调方法（FT、LoRA、DoRA）在查询矩阵（query matrices）上的更新幅度（ΔM）和更新方向（ΔD）在不同层（layers 1-6）和中间训练步骤（intermediate steps、final step）的变化情况。

• 子图 (a)：FT（全微调） 横轴为更新方向（ΔD），纵轴为更新幅度（ΔM）。可以看到，随着更新方向ΔD的增大，更新幅度ΔM呈下降趋势。不同颜色代表不同层，不同标记代表不同训练步骤，说明全微调时各层查询矩阵的更新幅度和方向存在一定的关联，且不同层、不同训练步骤的更新特征有差异。
• 子图 (b)：LoRA（低秩适应） 横轴ΔD增大时，纵轴ΔM呈明显上升趋势。这表明 LoRA 方法下，查询矩阵的更新方向变化与更新幅度变化呈正相关，且相比 FT，LoRA 的更新幅度整体更高（纵轴范围更大），说明 LoRA 在查询矩阵上的更新模式与全微调有显著区别，更倾向于大的幅度伴随大的方向变化。
• 子图 (c)：DoRA 横轴ΔD增大时，纵轴ΔM呈下降趋势，与 FT 的趋势类似但数值范围不同。这说明 DoRA 在查询矩阵的更新幅度和方向的关联模式上，和全微调有一定共性，但具体的数值表现（如幅度的大小范围）有所差异。

综上，这张图通过对比 FT、LoRA、DoRA 三种方法在查询矩阵更新上的幅度和方向关系，揭示了不同微调方法在模型参数更新模式上的差异，有助于理解它们在训练过程中对模型参数调整的不同特点。

DoRA 在微调时倾向于对预训练权重进行反向调整，其学习模式更接近全量微调，因此具备更优的学习能力。

4.2 DoRA 的梯度分析

由上面的公式可推导出损失函数 $L$ 对幅值 $m$ 和方向更新后向量 $V'$ 的梯度，具体如下： $$\nabla_{V'} L \propto \left( I - \frac{V' V'^T}{|V'|F^2} \right) \nabla{W'} L$$ $$\nabla_{V'} \mathcal{L} = \frac{m}{|V'|c} \left( I - \frac{V' V'^T}{|V'|c^2} \right) \nabla{W'} \mathcal{L}$$ 权重梯度 $\nabla{W'} L$ 会经过两个处理：一是被 $\frac{m}{|V'|_c}$ 缩放，二是被投影到远离当前权重矩阵的方向。这两个作用共同使梯度的协方差矩阵更接近单位矩阵，对优化过程是有利的。

$$\nabla_{m} \mathcal{L} = \frac{\nabla_{W'} \mathcal{L} \cdot V'}{|V'|_c}$$

创新点：通过权重分解，将幅值和方向调整解耦，让 LoRA 专注方向的低秩微调，同时放开幅值的可调性，从而提升微调的精细度和效率。

• DoRA 方法核心逻辑 可以把预训练模型的权重想象成'一支军队'：'幅值'是军队的'兵力强度'，'方向'是军队的'进攻路线'。之前的 LoRA 方法要同时调整'兵力强度'和'进攻路线'，就像一个指挥官同时指挥两项复杂任务，容易顾此失彼 —— 比如想微调路线时，兵力也会跟着按比例变化，没法做到'路线小改、兵力大调整'或反过来。 DoRA 的思路是'分工协作'：让 LoRA 专门负责调整'进攻路线'（方向），同时单独留出一个可调整的'兵力控制器'（幅值向量）。这样一来，两个任务互不干扰，指挥官能更精准地优化 —— 比如可以只把路线微调一点，同时大幅加强兵力；或者路线改得很多，兵力只微调。而且 DoRA 的'初始兵力和路线'是从预训练模型继承来的，不用从零开始，避免了'刚组建的军队没经验'的问题（即权重归一化的初始化敏感问题）。 另外，DoRA 在推理时能把调整后的'兵力和路线'重新合并成原来的权重格式，就像任务结束后军队恢复成标准编制，不会增加额外的'指挥成本'（推理延迟）。
• DoRA 学习模式的优势 全微调（FT）之所以效果好，是因为预训练模型已经有了'丰富的作战经验'（大量知识），面对新任务时，往往只需要'要么微调路线，要么调整兵力'，不用两者都大改 —— 比如面对小任务，可能只需要加强兵力，路线稍作调整就行。这种'一增一减'的模式，在数据上就表现为'方向差异（ΔD）和幅值差异（ΔM）负相关'（负斜率）。 而 LoRA 因为要同时调两者，只能做到'路线改多少，兵力也按比例改多少'（正相关，正斜率），灵活性差。DoRA 通过分工，也能实现 FT 那种'一增一减'的灵活调整：比如方向改得少（ΔD 小），幅值就能改得多（ΔM 大），反之亦然。从数据上看，DoRA 和 FT 的 ΔD 与 ΔM 都是负相关，而 LoRA 是强正相关，这说明 DoRA 学到了 FT 的'灵活作战思路'。
• 梯度分析的通俗理解 '梯度'可以理解为'指挥官的调整指令'——告诉军队'兵力该加多少、路线该改多少'。DoRA 的梯度计算有两个关键作用： 一是'校准指令强度'：把指令按当前的'兵力强度'（m）和'路线合理性'（$||V'||_c$）缩放，避免指令太激进或太保守，让调整更平稳； 二是'避免重复调整'：把指令投影到'远离当前路线'的方向，防止在同一个方向上反复微调，提升优化效率。 这种校准后的指令，会让'路线调整'（ΔV）的学习更稳定，不会像 LoRA 那样'路线和兵力乱同步'。而且从指令逻辑上能直接看出：路线改得少的场景，兵力调整指令会更强；路线改得多的场景，兵力调整指令会更弱 —— 这正好解释了 DoRA 为什么会呈现负斜率的学习模式。

总结

• 权重分解低秩适应（DoRA）：一种参数高效微调（PEFT）方法，核心是将预训练权重矩阵分解为幅值分量（列向量的范数，反映权重的'强度'）和方向分量（归一化后的列向量，反映权重的'特征捕捉方向'），并采用'分工优化'策略：通过 LoRA 对方向分量进行低秩微调（控制参数规模），同时将幅值分量设为独立可训练向量（提升调整灵活性）。其优势在于：规避权重归一化的初始化敏感性（从预训练权重初始化），实现与全微调（FT）相似的学习模式，且推理时可与预训练权重合并，无额外延迟。
• 幅值差异（ΔM）与方向差异（ΔD）：ΔM 衡量微调后权重与预训练权重在幅值分量上的平均差异；ΔD 衡量微调后权重与预训练权重在方向分量上的平均差异。两者的相关性是判断学习模式的关键指标：FT 与 DoRA 呈负相关（灵活调整，一增一减），LoRA 呈强正相关（同步比例调整，灵活性差）。
• 梯度投影与缩放：DoRA 梯度计算中的两个核心操作，用于优化方向分量（ΔV）的学习：缩放（通过 $\frac{m}{|V'|_c}$ 对权重梯度进行缩放）和投影（将梯度投影到方向分量的正交空间）。这两个操作共同使梯度协方差矩阵接近单位矩阵，提升优化稳定性与收敛速度。
• 学习模式的定量表征：通过 $(\Delta D, \Delta M)$ 的回归斜率与相关系数定量区分不同微调方法的学习特征：LoRA 为正斜率，FT 与 DoRA 为负斜率。DoRA 能够调整学习模式，摆脱 LoRA 的正斜率模式，更接近 FT 的学习模式。

4.3 训练开销的降低

1. 核心问题：在标准 LoRA 公式中，更新后权重 $W$ 与权重增量的梯度具有一致性；而 DoRA 将低秩适应定向于方向分量，导致低秩更新的梯度与 $W'$ 的梯度产生差异（如公式 6 所示），该差异在反向传播过程中引发额外内存开销。
2. 优化策略：在 DoRA 的训练开销优化中，需将 $|V + \Delta V|c$ 视为常数以降低内存消耗。视为常数，使其从梯度计算图中分离（即反向传播时不接收梯度）。此修改下，幅值分量 $m$ 的梯度保持不变，方向分量 $V'$ 的梯度被重新定义为：$\nabla_V \mathcal{L} = \frac{m}{C} \nabla{W'} \mathcal{L}$（其中 $C = |V'|_c$）。
3. 效果验证：通过对 LLaMA-7B 和 VL-BART 模型的消融实验表明，该优化可使 LLaMA 微调时训练内存降低约 24.4%，VL-BART 降低约 12.4%；且精度损失极小 ——VL-BART 精度完全不变，LLaMA 精度仅较未修改版本下降 0.2。
4. 后续应用：因该优化在内存节省与精度保持间实现高效平衡，后续所有 DoRA 相关实验均采用此调整策略。

5 实验（Experiments）

常识推理任务、单模态拓展到多模态、消融实验 结论：无论微调训练样本数量多少、秩的取值如何变化，DoRA 的性能均优于 LoRA。

5.1 常识推理

常识推理任务包含 8 个子任务，每个子任务都有预先定义好的训练集和测试集。将这 8 个子任务的训练集合并构建最终的训练集，并在每个子任务各自的测试集上进行评估。

为确保对比的公平性，研究人员最初按照 LoRA 的配置对 DoRA 模型进行微调：保持秩（rank）参数一致，仅调整学习率。

实验表格说明： 这张表格展示了不同大语言模型（LLaMA 7B/13B、LLaMA2 7B、LLaMA3 8B）在采用各种参数高效微调（PEFT）方法后，在八个常识推理数据集上的准确率对比，可从以下几个维度解读：

1. 模型与 PEFT 方法的性能差异

• 不同PEFT 方法（Prefix、Series、Parallel、LoRA、DoRA、DoRAⁱ）对同一模型的性能提升效果不同。例如，LLaMA-7B 采用 DoRA（Ours）时，在多个数据集（如 HellaSwag、ARC-e、OBQA 等）的表现优于 LoRA、Prefix 等方法，平均准确率（Avg.）达到 78.4，是该模型下 PEFT 方法中的最优表现。
• 不同模型规模的性能差距明显。如 LLaMA-13B 的整体表现优于 LLaMA-7B，LLaMA3-8B 在多数数据集上又优于 LLaMA2-7B，体现了模型规模对常识推理能力的影响；同时结合 PEFT 方法后，大模型的性能优势进一步凸显。

2. DoRA 方法的优势

• 无论是 DoRA 还是其调整版本 DoRAⁱ（秩减半），在多个模型（LLaMA-7B/13B、LLaMA2-7B、LLaMA3-8B）上的表现都较为突出，平均准确率往往高于 LoRA 等其他 PEFT 方法。例如，LLaMA3-8B 采用 DoRA（Ours）时平均准确率达到 85.2，是该模型组中 PEFT 方法的最优结果，甚至接近 ChatGPT 的 77.0（需注意 ChatGPT 是闭源大模型，架构与训练目标不同）。

3. 参数量占比与性能的权衡

• PEFT 方法的参数量占比# Params (%) 较低（如 DoRAⁱ普遍在 0.35-0.43% 左右），但能带来显著的性能提升，体现了参数高效微调在'少参数量消耗，高性能提升'上的优势。例如，LLaMA-7B 的 DoRAⁱ参数量占比仅 0.43%，平均准确率却达到 77.5，接近 ChatGPT 的 77.0，说明该方法在效率与效果上的平衡较好。

4. 数据集与任务的普适性

• 表格覆盖了多种常识推理任务数据集（BoolQ、PIQA、SIQA、HellaSwag、WinoGrande、ARC-e、ARC-c、OBQA），DoRA 在这些不同类型的常识推理任务上均表现出稳定的优势，说明其在常识推理场景下的普适性较强。

综上，这张表格清晰地对比了不同大语言模型在采用各类 PEFT 方法后的常识推理能力，突出了 DoRA 系列方法在参数效率和性能表现上的双重优势，同时也展示了模型规模、PEFT 方法对常识推理任务的影响，为选择常识推理场景下的大模型微调策略提供了数据支撑。

5.2 图像/视频 - 文本理解（多模态微调）

Table 2（图像 - 文本多任务评估）

• 参数量占比：LoRA 和 DoRA 的参数量占比仅约 5.9%，远低于全微调（FT）的 100%，体现了参数高效微调的'低参量消耗'优势。
• 性能表现：DoRA 的平均准确率（85.4）超过 LoRA 的 83.5，且在 TVQA（76.3）、How2QA（74.1）、TVC（45.8）、YC2C（145.4）的表现均优于 LoRA，部分任务（如 TVQA）甚至与全微调性能持平。

核心结论：在以 VL-BART 为骨干的多模态（图像 - 文本、视频 - 文本）任务中，DoRA 在参数量占比远低于全微调的前提下，性能表现优于 LoRA，且接近或达到全微调的水平，证明了 DoRA 在多模态参数高效微调场景下的有效性，实现了'低参量消耗'与'高性能表现'的平衡。

性能表现：DoRA 在平均准确率（Avg.）上达到 77.4，超过 LoRA 的 76.5，且接近全微调的 77.3；在具体任务中，DoRA 在 VQA（65.8）、GQA（54.7）、NVLR（73.1）、COCO Caption（115.9）的表现均优于 LoRA，部分指标接近甚至超过全微调。

Table 3（视频 - 文本多任务评估）

5.3 视觉指令微调

实验结论： 从表中可以观察到，LoRA 的平均准确率已经超过了全微调（FT），这可能意味着全微调（FT）存在过拟合问题。由于 DoRA 的设计初衷是提升 LoRA 的性能，使其更接近全微调（FT）的水平，因此在全微调（FT）性能不如 LoRA 的场景下，DoRA 相对 LoRA 的性能提升幅度，可能不会像在其他'全微调（FT）通常优于 LoRA'的实验中那样显著。尽管如此，DoRA 仍然表现仍优于 LoRA 和全微调（FT）：相较于 LoRA，平均性能提升 0.7%；相较于 FT，平均性能提升 1.1%。

5.4 DoRA 与其他 LoRA 变体的兼容性

W 的变化量通过不同的 LoRA 变体进行调整。对于 DoRA，引入的增量方向更新 V 的变化量这一概念同样可替换为其他 LoRA 变体。

DVoRA 融合了 DoRA 的优势特质，得分与 LoRA 相当甚至更高，但参数数量明显更少。 无论训练样本如何，DoRA 相较于 LoRA 和 VeRA 都能持续提升性能。

使用不同数量的 Alpaca 训练样本对 LLaMA2-7B 模型进行微调后，该模型在 MT-Bench 基准测试中的性能表现。 研究人员在图 3 中可视化了各方法在 LLaMA2-7B 模型上的平均性能，在附录的图 7 中可视化了各方法在 LLaMA-7B 模型上的平均性能。结果显示，在所有训练样本量设置下，DoRA 始终优于 LoRA，DVoRA 始终优于 VeRA。例如，当使用 7000 条训练样本时，DoRA 相较于 LoRA 的性能提升幅度为 0.3，DVoRA 相较于 VeRA 的性能提升幅度为 0.33；即便将样本量减少至 1000 条，DoRA 仍能以 0.29 的优势领先 LoRA，DVoRA 仍能以 0.22 的优势领先 VeRA。这表明，无论训练样本量大小，我们提出的方法（DoRA 和 DVoRA）相较于 LoRA 和 VeRA，均能持续提升模型性能。

5.5 DoRA 在不同秩设置下的鲁棒性

研究表明，在所有秩配置下，DoRA 的性能始终超过 LoRA。值得注意的是，当秩低于 8 时，性能差距会扩大，此时对于 ($r=8$)，LoRA 的平均准确率降至 40.74%；对于 ($r=4$)，其平均准确率为 39.49%。相比之下，对于 ($r=8$)，DoRA 保持着 77.96% 的显著准确率；对于 ($r=4$)，其准确率为 61.89%。这表明，无论秩设置如何，DoRA 都具有更强的稳健性，且性能始终优于 LoRA。

5.6 调优粒度分析

DoRA 可以仅通过更新特定模块的幅度和方向分量，来减少可训练参数的数量。

6 更广泛的影响（Broader Impacts）

6.1 QDoRA：对 QLoRA 的增强

在 Orca-Math 数据集（Mitra 等人，2024）上，LLaMA2-7B/LLaMA3-8B 模型分别采用 QDoRA、QLoRA 与全微调（FT）方法后的准确率对比。

为了进一步降低微调的内存需求，QDoRA 建议将预训练模型量化为 4 位，并在冻结的低位骨干网络之上对 LoRA 进行微调。 QDoRA 能够有效地将 QLoRA 的参数效率与全微调更精细的优化相结合——QDoRA 具有很大潜力，通过大幅降低微调大型语言模型的 GPU 内存需求给开源社区带来极大益处。

6.2 文本到图像生成

DoRA 相对于 LoRA 的优势能够延伸至图像生成任务中。

7 结论（Conclusion）

研究人员首先进行了一种新颖的权重分解分析，以揭示 LoRA 和 FT 之间不同的学习模式。基于这些见解，我们提出了 DoRA——一种与 LoRA 及其变体兼容的微调方法，其学习行为与 FT 更为相似。在各种微调任务和模型架构上，DoRA 的性能始终优于 LoRA。具体而言，DoRA 在常识推理和视觉指令调优任务上比 LoRA 有所改进。此外，在 Alpaca 指令调优任务上，DoRA 与 VeRA 也表现出兼容性。而且，DoRA 可以被视为 LoRA 的一种无成本替代方案，因为其分解后的幅度和方向分量在训练后可以合并回预训练权重中，确保没有额外的推理开销。对于未来的工作，希望能探索 DoRA 在语言和视觉之外领域的通用性，特别是在音频领域。

论文代码分析及创新点

class Linear(nn.Linear, LoraLayer):
    # Lora implemented in a dense layer
    def __init__(self, in_features: int,
                 out_features: int,
                 r: int = 0,
                 lora_alpha: int = 1,
                 lora_dropout: float = 0.0,
                 fan_in_fan_out: bool = False,
                 merge_weights: bool = True,
                 Wdecompose: bool = False,
                 dora_simple: bool = True,
                 **kwargs):
        nn.Linear.__init__(self, in_features, out_features, **kwargs)
        LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights)
        self.weight_m_wdecomp = nn.Linear(1, out_features, bias=False)
        self.fan_in_fan_out = fan_in_fan_out
        self.Wdecompose = Wdecompose
        self.dora_simple = dora_simple

        if self.Wdecompose == False:
            if r > 0:
                self.lora_A = nn.Linear(in_features, r, bias=False)
                self.lora_B = nn.Linear(r, out_features, bias=False)
            self.scaling = self.lora_alpha / self.r
            self.weight.requires_grad = False
            self.reset_parameters()
            if fan_in_fan_out:
                self.weight.data = self.weight.data.T

    def reset_parameters(self):
        nn.Linear.reset_parameters(self)
        if hasattr(self, "lora_A"):
            nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B.weight)

    def train(self, mode: bool = True):
        nn.Linear.train(self, mode)
        if self.Wdecompose == False:
            self.lora_A.train(mode)
            self.lora_B.train(mode)
            self.weight_m_wdecomp.train(mode)

    def eval(self):
        nn.Linear.eval(self)
        if self.Wdecompose == False:
            self.lora_A.eval()
            self.lora_B.eval()
            self.weight_m_wdecomp.eval()

    def forward(self, x: torch.Tensor):
        previous_dtype = self.weight.dtype
        if self.disable_adapters:
            raise NotImplementedError
        elif self.Wdecompose and not self.merged:
            norm_scale = self.weight_m_wdecomp.weight.view(-1) / (torch.linalg.norm(self.weight, dim=1))
            org_result = F.linear(x, transpose(self.weight, self.fan_in_fan_out))
            result = org_result + (norm_scale - 1) * (F.linear(self.lora_dropout(x), transpose(self.weight, self.fan_in_fan_out)))
            if not self.bias is None:
                result += self.bias.view(1, -1).expand_as(result)
        elif self.r > 0 and not self.merged:
            new_weight_v = self.weight + (self.lora_B.weight @ self.lora_A.weight) * self.scaling
            if self.dora_simple:
                norm_scale = self.weight_m_wdecomp.weight.view(-1) / (torch.linalg.norm(new_weight_v, dim=1)).detach()
            else:
                norm_scale = self.weight_m_wdecomp.weight.view(-1) / (torch.linalg.norm(new_weight_v, dim=1))
            org_result = F.linear(x, transpose(self.weight, self.fan_in_fan_out))
            dropout_x = self.lora_dropout(x)
            result = org_result + (norm_scale - 1) * (F.linear(dropout_x, transpose(self.weight, self.fan_in_fan_out)))
            if not self.bias is None:
                result += self.bias.view(1, -1).expand_as(result)
            result += (norm_scale * (self.lora_B(self.lora_A(dropout_x.to(self.lora_A.weight.dtype))))) * self.scaling
        else:
            result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)

        if result.dtype != previous_dtype:
            result = result.to(previous_dtype)
        return result