用法律合同微调Nemotron-8B

当我开始这个项目时，我想要的不仅仅是更好的分数。我想了解微调实际在做什么。本指南从基本概念开始，随着词汇出现进行解释，并使用我们真实的训练运行来展示什么有效，什么失败。

结果： 经过六次受控实验，验证Macro-F1从0.2956上升到0.8246。精确匹配准确率达到91.5%，200个示例验证集中的每个响应都可以解析。

这项工作是作为NVIDIA印度代理AI开放黑客马拉松的一部分构建的，我们的团队RTX Renegades构建了NemoCounsel，一个由本文描述的微调模型驱动的AI合同审查助手。感谢我的队友Aditya Rallapalli、Sainath Chakravadhanula、Ruchita Bhandari和Stuti Chourasia与我一起构建这个，以及我们的NVIDIA导师Raghul Asokan和Sutanshu Raj在整个过程中的指导。感谢NVIDIA提供使这些训练运行成为可能的GPU计算支持。

完整的实现是开源的，在NemoCounsel GitHub仓库中。该仓库包括Nemotron训练工作流、实验记录和重现运行所需的脚本。

1、问题

合同冗长、重复，包含大量含义取决于精确法律语言的条款。我们的任务是将每个输入条款映射到LEDGAR中的100个类别之一，LEDGAR是随LexGLUE法律语言基准分发的合同条款数据集。

我们将分类视为受限生成：模型看到一个指令、有效类别列表和条款文本，然后生成一个类别名称。这使得格式成为模型质量的一部分。语义上合理的答案如果无法映射到有效标签，仍然算作未解析。

基础模型是NVIDIA Llama-3.1-Nemotron-Nano-8B-v1。我们使用LoRA对其进行适配，并以4位NF4加载主干，遵循QLoRA推广的内存高效方法。结果是一个适配器，而不是另一个80亿参数模型的完整副本。

2、训练前了解LEDGAR

LEDGAR数据集具有令人愉悦的简单形状。每行包含text中的合同条款和label中的一个整数。该整数映射到100种条款类型之一，包括管辖法律、仲裁、保密、终止和弃权。仓库的LEDGAR概览笔记本以交互方式介绍相同的数据。

这是模型需要学习的那种映射：

条款：
"本协议应根据明尼苏达州法律解释和执行..."

标签：
管辖法律

这是一个分类任务，但我们使用了生成式语言模型。我们没有添加单独的100路分类头，而是用允许的标签提示模型，并训练它写出正确的标签名称。在评估期间，生成的文本必须映射回100个名称之一。

3、类别极不平衡

在60,000行的训练集中，管辖法律出现3,167次。书籍仅出现23次。最大类别大约是最小类别的137倍。因此，模型看到某些法律概念的频率远高于其他概念。

这不是一个小细节。模型可以通过学习频繁类别来提高整体准确率，同时继续在稀有类别上失败。我们需要一个评估指标，不让常见条款掩盖这些失败。

4、条款长度影响输入设计

在训练集的5,000个条款样本中，中位数条款包含85个单词，第95百分位数包含302个单词。最长的采样条款包含981个单词。大多数示例可以轻松放入一个模型输入，但少数仍可能超过固定的令牌限制。

这就是为什么截断后来变得重要。我们使用768个令牌的最大值，但正确标签附加在提示之后。如果我们不小心截断完整序列，可能会删除模型应该学习的答案。

5、用通俗英语解释微调

预训练语言模型已经从大量文本中学习了通用模式。微调不是从零开始教授语言。它给模型一个更窄行为的重复示例。在我们的案例中，该行为是：

指令 + 法律条款 → 正确的条款类别

在训练过程中，模型预测下一个令牌。训练程序将其预测与已知答案进行比较，并计算一个称为损失的数字。较大的损失意味着预测偏离目标更远。然后优化器进行小的参数更新，旨在减少下次类似的错误。在数千个示例上重复这个循环逐渐改变模型的行为。

有三个数据集需要在概念上分开：

训练数据产生参数更新。
验证数据不更新模型；它帮助我们比较实验。
测试数据应保持不变，直到配置冻结，给出对未见数据性能的最诚实估计。

我们使用了所有60,000个训练示例和LEDGAR 10,000行验证集中固定的200个示例样本。我们在每次实验后检查相同的样本。因此，最终数字是验证分数，而不是测试集结果。

6、为什么LoRA使8B模型变得实用

模型名称中的"8B"意味着大约80亿参数。更新和存储所有这些参数代价高昂。LoRA冻结原始权重，并在选定层旁边训练更小的低秩矩阵。将基础模型视为一本大教科书，LoRA适配器是一套紧凑的专家注释。在推理时，两者一起加载。

我们还以4位NF4加载冻结的主干，而不是将每个权重保持在16位或32位精度。量化大大降低了持有模型所需的内存。量化冻结权重与可训练LoRA适配器一起使用通常称为QLoRA。

四个术语将贯穿整个实验：

Epoch： 完整训练数据集的一次遍历。
批处理大小： 在累积更新之前一起处理的示例数量。
学习率： 每个优化步骤的大小。
LoRA秩（r）：可训练适配器的容量；更高的秩可以学习更复杂的更改，但使用更多内存且可能过拟合。

7、研究循环

该工作流借鉴了Andrej Karpathy的autoresearch的纪律：改变一个想法，运行整个评估，让指标决定该更改是否保留。

对于每个实验，我们：

写一个可测试的假设；
更改train.py并保持其描述有意义；
通过bash run_experiment.sh训练；
读取打印的验证Macro-F1；
仅当更改提高了最佳分数时才保留。

一些更改听起来合理，但仍然使结果更差。例如，将训练减少到一个epoch使Macro-F1从0.7548降低到0.6801。我们丢弃了该版本。

每次运行问一个受控问题：如果其他一切都保持不变，我们改变这一个选择，模型是否在相同的验证示例上有所改进？ 一次改变一件事帮助我们理解分数为何变动。

8、六次实验中的变化

最大的改进不是来自更大的秩或额外的epoch。它来自修复输入和目标构造。使用朴素截断，目标令牌可能被剪切，使模型没有答案可学习。为目标预留空间并仅截断提示在一次实验中使Macro-F1提高了0.4592。

该错误更容易用简化的训练序列理解：

[长提示和条款] [答案：管辖法律]

如果固定长度的令牌生成器简单地删除第768个令牌之后的所有内容，它可能会删除答案：管辖法律。示例仍然占用GPU时间，但模型不再收到它应该学习的部分。修复首先为答案预留空间，然后仅截断提示。这对任何初学者都是一个有用的教训：在调整超参数之前检查到达模型的确切令牌。

然后两个epoch在该保留检查点上恢复了另一个**+0.0393**。将适配器覆盖范围扩展到MLP投影适度帮助，将LoRA秩提高到32产生了最佳最终分数。

9、最终配置

主干使用NF4和bfloat16计算进行量化，使用Hugging Face文档化的bitsandbytes4位线性层。最终适配器针对q_proj、k_proj、v_proj、o_proj、gate_proj、up_proj和down_proj。PEFT参数直接对应于文档化的LoraConfig。

MODEL_NAME = "nvidia/Llama-3.1-Nemotron-Nano-8B-v1"
TRAIN_SIZE = 60_000
MAX_LENGTH = 768

LORA_R = 32
LORA_ALPHA = 64
LORA_DROPOUT = 0.05
LORA_TARGET_MODULES = [
    "q_proj", "k_proj", "v_proj", "o_proj",
    "gate_proj", "up_proj", "down_proj",
]

EPOCHS = 2
BATCH_SIZE = 4
GRAD_ACCUM_STEPS = 4
LEARNING_RATE = 2e-4

有效批处理大小为16。我们使用余弦学习率调度，10%预热，种子42，评估期间最多16个新生成的令牌。

10、损失曲线告诉我们什么

最终运行完成了两个epoch的7,500个优化器步骤。其Trainer摘要报告平均训练损失为0.1805，运行时间约为18,305秒（5小时05分钟）。该图表包含从最终原始日志中可恢复的每条损失记录，以及21点移动平均值；基础值发布为data/final_training_history.tsv。

损失对于检查优化很有用，但它不是选择指标。模型可以在不改进未见示例的情况下变得非常擅长其训练示例。这就是过拟合。训练损失告诉我们优化器是否从训练批次中学习。验证指标告诉我们这种学习是否转移。

11、为什么我们选择Macro-F1

准确率回答一个简单的问题：所有预测中有多少比例是正确的？这很直观，但它将每个条款视为一票。在不平衡的数据集中，大类别获得大多数投票。

考虑十个有三个标签的条款：

六个是管辖法律；
两个是终止；
两个是保密。

一个无用的分类器为每个条款预测管辖法律，得到六个正确答案。其准确率为60%，即使它从未识别出另外两种条款类型。

为了清楚地看到失败，我们为每个类别计算三个值：

精确率： 当模型预测此标签时，它正确的频率是多少？
召回率： 在所有真正具有此标签的条款中，它找到了多少？
F1： 精确率和召回率的调和平均值。只有当精确率和召回率都很高时，它才会变高。

Macro-F1为每个评估类别分别计算F1，然后取简单平均值：

Macro-F1 = (F1类别1 + F1类别2 + ... + F1类别N) / N

在我们评估笔记本中的工作示例中，一个不完美的分类器得分为0.80准确率和0.73 Macro-F1。懒惰的总是预测管辖法律的分类器仍然有0.60准确率，但其Macro-F1降至0.25。Macro-F1暴露了两个类别被忽略的事实。

Micro-F1在计算F1之前汇集所有预测，因此大类别仍然占主导地位。Weighted-F1对每个类别的分数取平均值，但按类别频率加权。两者都可以作为有用的辅助指标。对于LEDGAR，Macro-F1是主要指标，因为每个法律类别都应该平等影响结果。这也是LexGLUE基准中用于LEDGAR的主要指标。

Scikit-learn的F1文档精确定义了平均模式。在实践中，我们还记录了准确率和无法解析的生成答案数量。Macro-F1决定了实验是否保留。评估指标笔记本包含完整的工作示例，可以逐单元格运行。

12、重现实验

完整的代码和日志属于NemoCounsel仓库。确切的入口点是train.py、run_experiment.sh和all_attempts.tsv。为了忠实重现，请使用NVIDIA CUDA GPU并记录GPU型号、驱动程序、CUDA版本、Python版本、依赖版本和不可变的模型/数据集修订。

git clone https://github.com/RittikaJ/nemocounsel.git
cd nemocounsel/autoresearch/nemotron

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

python -c "import torch; print(torch.__version__, torch.version.cuda)"
python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"

bash run_experiment.sh

为了更强的出处，请固定Hugging Face资产，而不是依赖移动的默认值。模型卡公开了可提交寻址的文件，数据集加载器在其加载API中支持revision参数：

MODEL_REVISION = "485e5fa4b2ee7ff4a2bf78b8582251a49527f45d"
DATASET_REVISION = "fe9d5f8f254c32c0d9c343e86534523deb2fcb01"

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME, revision=MODEL_REVISION
)
dataset = load_dataset(
    "coastalcph/lex_glue",
    "ledgar",
    revision=DATASET_REVISION,
)

GPU上的可重现性不是一个绝对的保证。PyTorch明确指出，即使使用相同的种子，跨发布、平台或CPU/GPU执行也不保证相同的结果；其可重现性指南解释了确定性控制及其性能权衡。报告种子和完整环境。

13、限制和我们接下来会做什么

主要分数是200个示例的验证结果，而不是未接触的测试集结果。在选择假设时反复查阅了它，因此它可能高估了对真正未见合同的性能。下一个科学步骤是冻结最终配置，并在保留的测试集上精确评估一次。

在下一次运行中，我们会更改一个指标细节。Macro-F1在报告100类任务时应使用预期的固定标签集。将labels=range(100)传递给f1_score可以防止缺失类别静默更改平均值。我们还将报告每个类别的支持度和F1，以及置信区间。

原始最佳评分运行未保留其适配器权重。重新运行冻结配置可以创建适配器，但它应被描述为重现而不是确切的历史检查点。

14、我们学到了什么

最大的收益来自数据管道，而不是时尚的超参数。一旦目标在截断中幸存，模型就可以学习它应该被训练的任务。之后，更长的训练和更广泛、更高秩的适配器带来了更小的增量收益。

研究循环为我们提供了尝试的内容以及保留或拒绝每个更改的原因的记录。在六次Nemotron实验中，我们发现了一个主要的数据处理问题，拒绝了一个回归，并以0.8246验证Macro-F1配置完成。代码和数据可供任何想要重新运行或尝试下一个实验的人使用。

原文链接: I Fine-Tuned NVIDIA's 8-Billion-Parameter Language Model to Read Legal Contracts. Here's How.

汇智网翻译整理，转载请标明出处