LAMKIT: Length-Aware Multi-Kernel Transformers

• The paper introduces LAMKIT, highlighting its multi-kernel segmentation and length-aware encoding to counter context fragmentation and length overfitting in long documents.
• It integrates parallel multi-kernel encoding with segment-level positional embeddings and a pooling-fusion mechanism that preserves contextual coherency across various document scales.
• Empirical evaluations demonstrate that LAMKIT outperforms SOTA baselines with significant F1 improvements on health and legal document classification tasks.

Length-Aware Multi-Kernel Transformers (LAMKIT) constitute a hierarchical Transformer architecture designed to address context fragmentation and capacity overfitting endemic to long-document classification tasks. LAMKIT leverages multiple kernel granularities—distinct segmentations of the input—and an explicit length-aware vectorization module, resulting in robust performance across a wide spectrum of document lengths while maintaining competitive memory and computational efficiencies. LAMKIT’s design incorporates parallel multi-kernel encoding, segment-level positional embeddings, hierarchical integration of global document-length features, and a pooling-fusion mechanism for classification, yielding significant improvements over state-of-the-art (SOTA) baselines on health and legal text corpora (Han et al., 2024).

1. Architectural Overview

LAMKIT is composed of three principal components: (a) Multi-Kernel Encoding (MK), (b) Length-Aware Vectorization (LaV), and (c) Hierarchical Integration via document-level and length-level Transformers. The overall workflow can be summarized as follows:

• The input document $D$ (with $L$ tokens) is segmented at $K$ chunk sizes $m_1, ..., m_K$.
• Each segment is encoded by a shared pre-trained RoBERTa encoder, producing segment-level CLS vectors.
• Segment-position embeddings, derived from a sine-cosine function at the segment (not token) level, are added to these vectors.
• Pooling over segment representations produces length vectors for each kernel, which are then processed by a dedicated "length encoder" Transformer, producing $K$ global length-aware codes.
• Each kernel’s sequence of segment representations is processed by a two-layer "document encoder" Transformer. The resulting document encodings are summed with the broadcasted length code across all segments of each kernel.
• Max and average pooling are performed per kernel; their concatenations across kernels are averaged to yield a single document-level vector.
• This embedding is passed through a classification head (linear + softmax/sigmoid) for final prediction.

Key architectural parameters include three kernel sizes (e.g., {128, 256, 512} tokens for MIMIC and SCOTUS; {32, 64, 128} for ECtHR), RoBERTa-base as backbone, Transformer layers (length encoder: 1 layer; document encoder: 2 layers, 12 heads, $d = 768$), and non-overlapping segmentation.

2. Mathematical Formulation

The model instantiates multi-head self-attention and segment-level position embeddings as follows:

• Segment Position Embedding:

For segment index $j$ and dimension $i$:

$$\mathrm{PE}_{(j,2i)} = \sin \left( \frac{j}{10000^{2i/d}} \right),\quad \mathrm{PE}_{(j,2i+1)} = \cos \left( \frac{j}{10000^{2i/d}} \right)$$

Added to each RoBERTa-encoded CLS:

$x_{i,j} = h_{i,j}^{seg} + \mathrm{PE}^{seg}(j)$

• Multi-head Attention:

$L$0

• Length Vector:

For kernel $L$1, pooling across $L$2 segments:

$L$3

$L$4 form the input for the length encoder Transformer.

• Hierarchical Pooling:

After integration:

$L$5

Concatenate $L$6 per kernel, and finally average after concatenation across all $L$7 kernels.

The final document vector is thus sensitive both to multi-scale context and global length features.

3. Addressing Context Fragmentation and Length-Overfitting

LAMKIT’s multi-kernel segmentation ensures that contiguous context is preserved at multiple scales—reducing the likelihood that crucial spans (such as sentence boundaries or well-formed linguistic units) are consistently split across all granularities. This multi-scale redundancy mitigates context fragmentation, a challenge with fixed-chunk models in which a salient phrase may be split at every boundary.

Length-overfitting arises in models trained primarily on documents of specific lengths; such models exhibit a drop in classification efficacy when faced with much shorter or longer documents. LAMKIT introduces an explicit length-encoding pathway: global document-length characteristics, distilled through the length encoder, are summed into each per-segment representation. This mechanism increases the representational robustness of the network across a wide span of document lengths, empirically yielding stability for documents ranging from several hundred to over ten thousand tokens.

Regularization is imposed using dropout ($L$8), AdamW optimizer (weight decay $L$9), early stopping (patience $K$0 on F1-micro), and mixed-precision (fp16) training.

4. Empirical Evaluation and Benchmarks

LAMKIT was benchmarked on five long-document classification tasks from health and legal domains:

Dataset	Avg Length ($K$1)	Train Size	Labels	Domain
Diabetes	720	1,265	10 (binary)	Clinical notes (health)
MIMIC-III	2,200	11,368	50 (multi)	ICU discharge summaries
ECtHR-A	2,139	11,000	11	Court cases (law)
ECtHR-B	2,139	11,000	11	Paragraphed court cases
SCOTUS	9,840	7,800	14	US Supreme Court opinions

Baselines included BERT (truncated), Longformer, BigBird, and hierarchical single-kernel Transformers.

LAMKIT achieved the following improvements over the mean performance of all baselines:

• F1-micro: +4.7% absolute gain
• F1-macro: +7.2% absolute gain
• On MIMIC-III: +8.0% F1-micro, +10.9% F1-macro

Quartile-based robustness analysis sorted documents by length and showed that LAMKIT maintains stable improvements (up to +8.6% absolute in some quartiles), in contrast to baseline models which can exhibit overfitting to specific length bands.

5. Ablation Studies and Model Robustness

Ablation analysis confirmed LAMKIT’s reliance on both its multi-kernel structure and explicit length-aware vectorization. Removing either component resulted in average F1-micro drops of 1.3%, and up to 2.8% if both were ablated. F1-macro performance suffered even more acutely (drops of 1.9% to 3.5%), underscoring the mutual necessity of multi-granular encoding and global length conditioning.

Configuration	Δ F1-micro	Δ F1-macro
- Multi-Kernel only	-1.3%	-1.9%
- Length-Aware only	-1.3%	-2.4%
- Both (MK & LaV) ablated	-2.8%	-3.5%

A plausible implication is that single-kernel hierarchical models cannot by themselves offer robustness to variations in document length or boundary integrity, and length awareness must be injected explicitly.

6. Implementation Details and Training Regimen

LAMKIT architecture can be instantiated using the provided pseudocode, which details the multi-kernel segmentation, segment embedding, length pooling, and hierarchical Transformer stacking. Each kernel uses a RoBERTa-base encoder (12 layers, $K$2), with two Transformer layers per document encoder and a single layer per length encoder.

Training used a batch size of 16 or 32, learning rate in the range $K$3, AdamW optimizer with weight decay, up to 20 epochs, and early stopping. Mixed-precision (fp16) was employed for computational efficiency. Strides equaled kernel size; segments were non-overlapping.

7. Comparative Context and Related Work

LAMKIT advances over SOTA by jointly addressing context fragmentation and length-overfitting, which distinguished it from prior models such as Longformer, BigBird, and H-BERT—all of which either operate at a single scale or lack mechanisms for explicit document length conditioning.

In the broader arena of length-extrapolation in transformers, methods such as MEP ("MEP: Multiple Kernel Learning Enhancing Relative Positional Encoding Length Extrapolation" (Gao, 2024)) consider multi-kernel bias in relative positional encoding to extend sequence generalization in self-attention, though they focus on the architectural manipulation of attention scores via kernel mixtures (e.g., exponential, Gaussian) and do not directly tackle document segmentation or global representation pooling as in LAMKIT. This suggests that LAMKIT and MEP represent complementary directions: the former in hierarchical, length-aware document modeling; the latter in kernelized positional encoding within single-pass attention architectures.

LAMKIT’s empirical and architectural contributions establish a new standard for long-document classification, demonstrating that multi-kernel, length-conditioned hierarchical Transformers can robustly generalize across both context boundaries and document lengths (Han et al., 2024).