Enhancing BiGRU with a KAN Block for Legal Document Classification and Summarization

Technical Summary: Enhancing BiGRU with a KAN Block for Legal Document Classification and Summarization

Problem Statement
The paper addresses the challenges of classifying and summarizing legal documents within a low-resource, multilingual setting. Specifically, it targets legal texts from Bangladesh, which present a unique combination of difficulties:

• Multilingualism: The dataset contains a mix of Bengali, English, and transliterated (Romanized) Bengali, creating significant vocabulary and syntactic variations.
• Domain Specificity: Legal documents (case notes, rulings, judgments) are lengthy, structured, and contain specialized terminologies that require capturing both local semantics and long-term dependencies.
• Class Imbalance: The target variable, "Disposition," is divided into 10 classes with highly uneven distributions, leading to bias toward majority classes in standard training.
• Resource Constraints: The study operates under limited computational resources, necessitating efficient architectures rather than massive pre-trained models.

Existing approaches, including classical machine learning (SVM, Random Forest) and standard deep learning models (BiLSTM, BiGRU), struggle with these specific constraints, particularly in modeling intricate non-linear relations and handling the multilingual noise without massive parameter counts.

Methodology
The authors propose a novel architecture that augments existing recurrent networks with Kolmogorov-Arnold Networks (KANs). Instead of replacing the recurrent backbone, the KAN block is used as an enhancement module to improve representation learning.

1. Dataset: The study utilizes a dataset of 2,937 legal documents sourced from Manupatra. The data is split into 2,349 training cases and 588 held-out evaluation cases. Preprocessing includes handling missing values, text normalization, and tokenization. To address class imbalance, a WeightedRandomSampler is employed during training.
2. Classification Architecture:
   • Backbone: A Bidirectional Gated Recurrent Unit (BiGRU) processes the input token sequence to capture bidirectional context.
   • Aggregation: Mean pooling and max pooling are applied to the hidden states to generate a fixed-size document representation ($h_{doc}$).
   • KAN Enhancement: The document representation is passed through a KAN block ($\hat{h}_{doc} = KAN(h_{doc})$) before the final classification head. The KAN block introduces parameterized spline-based edge functions to learn complex non-linear transformations of the recurrent features.
3. Summarization Architecture:
   • Backbone: An attention-based GRU encoder-decoder framework.
   • KAN Enhancement: A KAN head is added above the attention-based recurrent model to refine the representation before token prediction, enhancing the model's expressive capability.
4. Training: Models are trained for 200 epochs using the Adam optimizer, a learning rate of $2 \times 10^{-5}$, and a batch size of 8.

Key Contributions

• Novel Architecture: Introduction of a BiGRU-KAN model for legal classification and an attention-based GRU-KAN model for legal summarization.
• Multilingual Low-Resource Focus: Successful application of these models on a challenging Bangladeshi legal dataset containing Bengali, English, and Romanized Bengali.
• Ablation Validation: Demonstration that the KAN block is a critical component, significantly boosting performance over the baseline recurrent models without KAN.
• Imbalance Mitigation: Implementation of WeightedRandomSampler to handle the skewed distribution of disposition labels.

Experimental Results
The proposed models were compared against classical machine learning algorithms (Logistic Regression, SVM, Random Forest, etc.) and Pre-trained Language Models (PLMs) like BERT, RoBERTa, and T5, as well as recurrent baselines (BiLSTM, BiGRU without KAN).

• Classification Performance:
  • The BiGRU + KAN model achieved an accuracy of 67.96% and a weighted F1 score of 0.65.
  • This outperformed the best classical ML models (Random Forest and SVM at 62%) and significantly surpassed the BiGRU baseline without KAN (57.34%).
  • Notably, the proposed model outperformed several PLMs (e.g., Legal-BERT at 38.85%, RoBERTa at 37.41%), which the authors attribute to the constrained resources and hyperparameter tuning differences during the comparison.
• Summarization Performance:
  • The AttnGRU + KAN model achieved ROUGE F1 scores of 0.38 (ROUGE-1), 0.23 (ROUGE-2), and 0.31 (ROUGE-L).
  • These results were superior to BiLSTM (0.30/0.18/0.25) and Pointer-Generator networks (0.35/0.20/0.28).
• Ablation Study: The removal of the KAN block resulted in a drop in classification accuracy from 67.96% to 57.34%, confirming the block's efficacy in enhancing non-linear feature transformation.

Significance and Claims
The paper claims that integrating KAN blocks into recurrent architectures is a viable and effective strategy for legal NLP tasks, particularly in low-resource, multilingual environments. The study suggests that:

1. KANs as Enhancers: KANs serve as powerful architectural improvements that refine the representations learned by standard RNNs, offering better class separation and handling of non-linear legal nuances.
2. Efficiency: The proposed approach achieves superior results compared to both classical ML and several PLMs under resource-constrained conditions, suggesting that architectural innovation can sometimes outweigh the need for massive pre-training in specific domains.
3. Limitations: The authors modestly acknowledge remaining challenges, including the difficulty of predicting minority classes despite sampling techniques, the complexity introduced by the multilingual nature of the data, and the occasional omission of procedural details in summaries. They note that comparisons with PLMs were limited by computational constraints and tuning budgets.

In conclusion, the study posits that architectural enhancements, specifically the addition of KAN blocks to base RNNs, are crucial for advancing performance in resource-constrained legal NLP applications.