PaperlayerView exact revision on arXiv

Source-linked AI summary

Exploring Selective Retrieval-Augmentation for Long-Tail Legal Text Classification

Boheng Mao

arXiv:2508.19997v3cs.CLcs.IR

TL;DR

Long-tail legal classification leaves rare labels difficult to predict, motivating a selective alternative to broad retrieval augmentation. The paper uses SRA to retrieve training-set examples for low-frequency classes without changing the classifier architecture. Across LEDGAR and UNFAIR-ToS, SRA consistently improves micro-F1 and macro-F1 over LexGLUE baselines, although the proof-of-concept relies on ground-truth labels at validation and test time.

  • Problem

    Long-tail legal datasets underrepresent many labels, and models struggle on rare categories.

    Source: §1 · p. 1

  • Method

    SRA selectively augments low-frequency classes with semantically relevant examples retrieved from the training set, without changing the classifier architecture.

    Source: §1 · pp. 1–2

  • Results

    SRA consistently improves both micro-F1 and macro-F1 over LexGLUE baselines on LEDGAR and UNFAIR-ToS.

    Source: §1 · p. 2, §6 · p. 6

  • Takeaways & Limitations

    Tail-aware augmentation shows potential for addressing long-tail distributions in legal NLP.

    Source: §6 · p. 6

    The evaluation uses ground-truth labels during validation and testing to identify low-frequency cases, which is unrealistic at deployment.

    Source: Limitations · p. 6

Abstract

from arXiv · show

Legal text classification is a fundamental NLP task in the legal domain. Benchmark datasets in this area often exhibit a long-tail label distribution, where many labels are underrepresented, leading to poor model performance on rare classes. This paper explores Selective Retrieval-Augmentation (SRA) as a proof-of-concept approach to this problem. SRA focuses on augmenting samples belonging to low-frequency labels in the training set, preventing the introduction of noise for well-represented classes, and requires no changes to the model architecture. Retrieval is performed only from the training data to ensure there is no potential information leakage, removing the need for external corpora simultaneously. SRA is tested on two legal text classification benchmark datasets with long-tail distributions: LEDGAR (single-label) and UNFAIR-ToS (multi-label). Results show that SRA achieves consistent gains in both micro-F1 and macro-F1 over LexGLUE baselines.

1 Introduction

Legal text classification is challenged by long-tail label distributions, which make rare categories difficult to classify. The paper introduces SRA, which selectively augments low-frequency classes using in-domain retrieval while keeping the classifier architecture unchanged.

  • Long-tail label distributions cause models to perform well on common categories but struggle on rare ones, reducing overall performance.Source: §1 · p. 1
  • Retrieval-augmented methods often augment all samples or rely on external corpora, motivating a simpler legal-domain approach focused on underrepresented labels.Source: §1 · p. 1
  • SRA selectively augments low-frequency legal-text classes with semantically relevant examples retrieved entirely from the training set.The approach targets underrepresented labels rather than applying retrieval to every sample.Source: §1 · pp. 1–2
  • SRA uses in-domain legal retrieval without external knowledge and leaves the classifier architecture unchanged.The evaluation assumes access to ground-truth labels to isolate the effect of selectively augmenting low-frequency classes.Source: §1 · pp. 1–2
  • SRA achieves consistent gains over LexGLUE baselines across LEDGAR and UNFAIR-ToS in both micro-F1 and macro-F1.The evaluation covers LEDGAR as a single-label benchmark and UNFAIR-ToS as a multi-label benchmark.Source: Figure 1 · p. 1, §1 · p. 2

2 Related Work

Prior work addresses long-tail legal classification through loss- or classifier-level methods and retrieval augmentation, but these lines rarely combine contextual support for rare labels with in-domain retrieval. SRA positions selective retrieval from training data as a complementary approach.

2.1 Legal Benchmarks and Long-tail Classification

  • Long-tail methods such as logit adjustment and decoupled representation learning alleviate imbalance at the loss or classifier level but do not directly provide contextual support for infrequent classes.Source: §2.1 · p. 2
  • SRA complements existing imbalance methods by selectively incorporating in-domain examples for underrepresented labels.Source: §2.1 · p. 2

2.2 Retrieval-Augmented Classification

  • Retrieval augmentation has been extended from knowledge-intensive external retrieval to multi-label, few-shot, and zero-shot classification settings.Source: §2.2 · p. 2
  • Earlier retrieval approaches rarely focus on underrepresented labels in skewed datasets and often target open-domain settings requiring external information.Source: §2.2 · p. 2
  • SRA builds on retrieval-augmented classification by retrieving within the training set to support less common legal categories without outside sources.Source: §2.2 · p. 2

3 Methodology

SRA identifies low-frequency labels, retrieves semantically related training examples for matching samples, and concatenates them with the original text. The method fixes the low-frequency set from training data, controls sequence length, and preserves the standard RoBERTa architecture.

3.1 Selective Retrieval

  • SRA defines low-frequency classes by ranking training-set labels by frequency and selecting the bottom α proportion.The cutoff α is predefined and lies in (0, 1].Source: §3.1 · p. 2
  • For samples assigned to low-frequency classes, SRA selects candidates with TF-IDF and reranks them using Sentence-BERT cosine similarity, excluding self-retrieval.The default number of retrieved candidates is k=1.Source: §3.1 · p. 3
  • Retrieved clauses are concatenated with the original text using prompts that distinguish the original clause from related reference clauses.Source: §3.1 · p. 3

3.2 Augmentation Strategy

  • The low-frequency label set is derived from training data, fixed across splits, and paired with training-only retrieval to prevent information leakage.Source: §3.2 · p. 3
  • Retrieved clauses are truncated to 64 tokens and sequences are capped at 512 tokens for RoBERTa-base, while most augmented samples remain within these limits.In LEDGAR, 95% of augmented samples contain fewer than 421 tokens; truncation rarely occurs in UNFAIR-ToS.Source: Figure 2 · p. 2, §3.2 · p. 3
  • SRA is a data-level strategy that uses the unchanged RoBERTa-base encoder and standard linear classification layer.This preserves compatibility with usual training procedures.Source: §3.2 · p. 3

4 Experiments

Experiments evaluate SRA on two LexGLUE benchmarks with different long-tail patterns, comparing selective augmentation against standard and full-retrieval baselines. SRA improves both datasets, while ablations show that effective selection depends on the underlying imbalance.

4.1 Datasets

  • LEDGAR contains 100 single-label contract-clause categories with highly dispersed frequencies, whereas UNFAIR-ToS is multi-label and dominated by empty-label samples.The datasets therefore instantiate different forms of long-tail imbalance.Source: §4.1 · p. 3

4.2 Implementation Details

  • Macro-F1 is the primary metric because it better captures performance on uncommon classes under long-tail distributions.Micro-F1 is also reported as an overall performance measure across classes.Source: §4.2 · p. 4

4.3 Baselines

  • SRA augments only low-frequency classes, while the baseline uses no retrieval and full retrieval augments every sample.The experiments also include best LexGLUE models for comparison, with augmentation cutoffs selected on validation data.Source: §4.3 · p. 4, §4.4 · p. 4

4.4 Main Results

  • 0.931 / 0.887 on LEDGAR and 0.988 / 0.924 on UNFAIR-ToS are the best SRA micro-F1 / macro-F1 results, surpassing baseline and LexGLUE comparisons.Full retrieval lowers macro-F1 on both datasets, whereas SRA targets low-frequency classes or non-empty labels.Source: §4.4 · pp. 4–5

4.5 Ablation Studies

  • LEDGAR macro-F1 is maximized between 55% and 65% augmentation, after which performance drops as augmentation becomes too extensive.This supports selective augmentation rather than augmenting all classes.Source: §4.5 · p. 5, Figure 3 · p. 5
  • UNFAIR-ToS performs best when SRA selects all non-empty labels, while augmenting all samples lowers macro-F1 because its imbalance is concentrated in empty-label examples.Thus, the optimal selective strategy follows the dataset’s imbalance structure.Source: §4.5 · p. 5

5 Analysis

Analysis tests whether SRA’s gains are statistically reliable, concentrated in long-tail frequency buckets, and explained by retrieval coverage rather than changing retrieval quality. The results show improvements for rare and medium-frequency labels as well as selected individual categories.

Statistical significance

  • A +0.060 LEDGAR macro-F1 gain is supported by a 95% bootstrap CI of [0.0459, 0.0662] and McNemar’s test with p ≪0.01.The analysis uses the 65% cutoff.Source: §5 · p. 5

Bucketed performance

  • Macro-F1 improves from 0.759 → 0.808 for rare labels and 0.827 → 0.905 for medium-frequency labels, while frequent labels improve from 0.904 → 0.948.The largest relative gains occur in rare and medium-frequency buckets without sacrificing common-class performance.Source: §5 · p. 5

Top-k categories

  • Label 4 improves from 0.34 → 0.76, while several very low-frequency categories gain more than +0.20 macro-F1.These label-wise results show substantial improvements for selected underperforming classes.Source: §5 · pp. 5–6

Coverage and similarity

  • As the cutoff increases, retrieval coverage rises while cosine similarity remains stable, indicating that augmentation quantity rather than retrieval quality drives performance differences.Figure 4 examines these trends on LEDGAR.Source: §5 · p. 6, Figure 4 · p. 6

UNFAIR-ToS analysis

  • On UNFAIR-ToS, selective augmentation of all non-empty instances reaches 0.988 micro-F1 and 0.924 macro-F1 without harming the empty class.This extends the tail-aware pattern to a different imbalance structure.Source: §5 · p. 6

6 Conclusion

The paper introduces SRA as an exploratory approach that directs retrieval augmentation toward tail classes in legal text classification. Results indicate substantial gains over uniform augmentation and strong LexGLUE baselines, motivating further study of tail-aware augmentation.

6 Conclusion

  • SRA directs augmentation toward tail classes and improves both micro-F1 and macro-F1 over uniform augmentation and strong LexGLUE baselines.The paper characterizes the implementation as exploratory rather than definitive.Source: §6 · p. 6

Limitations

SRA’s evaluation assumes ground-truth labels during validation and testing to identify low-frequency classes, which is unrealistic at deployment.

  • Ground-truth labels determine whether validation and test samples belong to the low-frequency set used for selective augmentation.The low-frequency set is defined from the training distribution and applied during evaluation.Source: Limitations · p. 6
  • Because labels are unavailable at inference time, this setup serves only as a proof of concept rather than a deployment-ready procedure.The paper suggests uncertainty-based criteria or auxiliary models as possible labelfree alternatives.Source: Limitations · p. 6

Ethics Statement

The study uses public legal datasets without personal or sensitive information, while cautioning that dataset biases may affect outputs in real-world legal decisionmaking.

Ethics Statement

  • The research uses publicly accessible LEDGAR and UNFAIR-ToS datasets and includes no personal or sensitive information.Source: Ethics Statement · p. 6
  • Potential dataset biases may be reflected in model outputs, so applying the models to real-world legal decisionmaking requires care.Source: Ethics Statement · p. 6

A Dataset Statistics

The datasets exhibit highly imbalanced long-tail label distributions, and the retrieval pipeline combines sparse TF-IDF retrieval with dense SBERT reranking to select one training-set clause.

A Dataset Statistics

B Retrieval Settings

  • All train, validation, and test splits are processed with the same retrieval settings.Source: B Retrieval Settings · p. 7
  • The retrieval pipeline uses word-level TF-IDF unigrams and bigrams to retrieve the top 20 candidates for each query.Terms appearing in fewer than two documents or more than 80% of the corpus are filtered, with sublinear term-frequency scaling.Source: B Retrieval Settings · p. 7
  • SBERT reranks the 20 TF-IDF candidates by cosine similarity, retains the top five, and uses only the best clause for augmentation.The dense encoder uses normalized embeddings and truncates each retrieved clause to 64 tokens.Source: B Retrieval Settings · p. 7

C Additional Results on LEDGAR

On LEDGAR, the 65% cutoff produces the largest reported macro-F1 gain, while lower and higher cutoffs yield smaller gains.

C Additional Results on LEDGAR

  • The LEDGAR analysis groups classes by frequency, and the UNFAIR-ToS distribution separately distinguishes empty from non-empty labels.Source: Table 3 · p. 8, Table 4 · p. 8
  • +0.0389 macro-F1 gain occurs at the 20% cutoff, with a 95% CI of [0.0285, 0.0493] and p < 10^-12.The largest bucket-level gain is reported for low-frequency classes at +0.0915.Source: C Additional Results on LEDGAR · p. 8
  • +0.0566 macro-F1 gain occurs at the 65% cutoff, with a 95% CI of [0.0459, 0.0662] and p < 10^-82.Both mid- and high-frequency buckets improve significantly, while low-frequency classes improve moderately.Source: C Additional Results on LEDGAR · p. 8
  • +0.0211 macro-F1 gain occurs at the 90% cutoff, with a 95% CI of [0.0126, 0.0294] and p < 10^-21.All frequency buckets improve slightly under this cutoff.Source: C Additional Results on LEDGAR · p. 8

D Case Study

The case studies show that SRA can either correct baseline errors or introduce misleading signals. In LEDGAR, retrieval emphasized positional cues over duties, while in UNFAIR-ToS, removal-related cues enabled the correct prediction.

Case Study

  • Retrieval can improve predictions but can also introduce misleading signals, as illustrated by one LEDGAR failure and one UNFAIR-ToS success.These cases provide representative examples of both outcomes rather than establishing that retrieval is uniformly beneficial for every instance.Source: D Case Study · p. 8

Case 1 (LEDGAR)

  • SRA misclassified an LEDGAR clause as Positions (Class 69) instead of the gold label Duties (Class 32) after retrieval emphasized the executive’s title.The baseline correctly predicted Duties (Class 32), but retrieved references to the “President and CEO” introduced positional cues that overrode duty-related context.Source: D Case Study · p. 8

Case 2 (UNFAIR-ToS)

  • SRA corrected an UNFAIR-ToS baseline error by predicting Content removal (Class 3) from a retrieved clause containing removal and infringement cues.The baseline predicted no label, whereas the retrieved clause reinforced cues such as “remove/delete content,” “infringing,” and “repeat infringers.”Source: D Case Study · p. 8

Related papers

Loading 2508.19997v3…