This experiment fine-tunes Meta's LLaMA 3.1 8B model for legal document classification—a task requiring nuanced understanding of contract clauses, regulatory language, and legal terminology. Our goal is to create a specialized model capable of categorizing legal documents into 23 distinct classes with >85% accuracy, replacing our current rule-based system that achieves only 67% accuracy.
We're using LoRA (Low-Rank Adaptation) to efficiently fine-tune the model while keeping computational costs manageable. This approach allows us to train adapter weights (~16M parameters) rather than the full 8B parameter model, reducing GPU memory requirements by approximately 75% while maintaining comparable performance to full fine-tuning.
Why Legal Document Classification?
Legal teams spend an estimated 60% of their time on document review and classification—a tedious but critical task that directly impacts contract management, compliance monitoring, and risk assessment. Our client processes over 50,000 legal documents monthly, and even small improvements in classification accuracy translate to significant time savings and reduced compliance risk.
Previous attempts using BERT-based classifiers achieved 78% accuracy but struggled with long documents and nuanced legal language. The hypothesis for this experiment is that a larger language model with domain-specific fine-tuning can better capture the contextual relationships in legal text, particularly for ambiguous clauses that require reasoning across multiple paragraphs.
The Legal-Docs-V2 Dataset
We curated a dataset of 45,000 legal document samples across 23 categories, including NDAs, employment contracts, licensing agreements, terms of service, privacy policies, and various regulatory filings. Each document was reviewed by at least two legal professionals, with disagreements resolved by a senior attorney. The inter-annotator agreement (Cohen's Kappa) was 0.87, indicating high consistency.
Data Split & Preprocessing
- Truncated to 4096 tokens max
- PII redaction applied
- Stratified sampling by class
The dataset has significant class imbalance (NDAs represent 18% while regulatory filings are only 2%). We address this through weighted cross-entropy loss during training, with class weights inversely proportional to frequency. We also use oversampling for minority classes during data loading.
Loss Curves
GPU Cluster Status
All OnlineTraining Analysis
The training curves show healthy convergence with the train-validation gap remaining relatively stable after epoch 10. The initial rapid decrease in loss (epochs 1-8) corresponds to the model learning basic legal terminology and document structure. The slower convergence phase (epochs 9-18) represents fine-grained learning of nuanced classification boundaries.
We observed early signs of overfitting around epoch 12, where validation loss began to plateau while training loss continued to decrease. This prompted us to increase dropout from 0.05 to 0.1 and switch from a linear learning rate schedule to cosine annealing, which helped stabilize the validation loss.
We track four primary metrics during training, evaluated on the held-out validation set at each epoch. All metrics are currently exceeding our target thresholds, suggesting the model is on track to meet or exceed our accuracy goals.
Per-Class Performance (Top 5 & Bottom 3)
Lower-performing classes correlate with smaller training sample sizes. Post-training, we plan to collect additional samples for these categories and perform targeted fine-tuning.
Why LoRA Over Full Fine-Tuning?
Low-Rank Adaptation (LoRA) freezes the pre-trained model weights and injects trainable rank decomposition matrices into each layer of the Transformer architecture. For our 8B parameter model, this reduces trainable parameters from 8 billion to approximately 16 million—a 500x reduction that dramatically decreases memory requirements and training time.
Full Fine-Tuning
- ~160GB GPU memory required
- 8B parameters to train
- Catastrophic forgetting risk
- ~$2,400 compute cost
LoRA Fine-Tuning
- ~40GB GPU memory required
- 16M parameters to train
- Preserves base capabilities
- ~$580 compute cost
Hyperparameters
Training Log
Last 5 epochs| Epoch | Train Loss | Val Loss | LR | Time |
|---|---|---|---|---|
| 14 | 0.5400 | 0.8600 | 1.2e-5 | 12m 34s |
| 15 | 0.5200 | 0.8500 | 1.0e-5 | 12m 28s |
| 16 | 0.5000 | 0.8400 | 8.0e-6 | 12m 31s |
| 17 | 0.4800 | 0.8400 | 6.0e-6 | 12m 29s |
| 18 | 0.4700 | 0.8300 | 4.0e-6 | 12m 35s |
Epoch 15: Switched to cosine annealing learning rate schedule. Validation loss has stabilized and we're seeing consistent improvements in F1 score across minority classes.
Epoch 12: Early signs of overfitting detected (train-val gap widening). Increased dropout from 0.05 to 0.1 and added gradient clipping at 1.0. This stabilized training.
Epoch 8: LoRA adapters showing good convergence. Baseline accuracy of 85% achieved ahead of schedule. Considering early stopping if validation loss doesn't improve by epoch 20.
Epoch 5: Initial experiments with rank=8 showed underfitting. Increased to rank=16 which improved capacity without significant memory overhead.
What Happens After Training Completes
Once training completes (~7 more epochs), we'll run comprehensive evaluation on the held-out test set, perform error analysis on misclassified samples, and conduct A/B testing against the current production system before full deployment.
Immediate Next Steps
- • Run full test set evaluation
- • Generate confusion matrix analysis
- • Export model for inference optimization
- • Benchmark inference latency
Deployment Preparation
- • Quantize to INT8 for production
- • Set up shadow deployment pipeline
- • Configure monitoring dashboards
- • Plan gradual traffic migration