lalala512/lexlattice-2d-nca

LexLattice: 2D-NCA semantic-lattice consolidator for EUR-Lex-Sum

LexLattice is a lightweight (1.8M-parameter) multilingual extractive summarizer for long EU legal acts, covering all 24 official EU languages of the EUR-Lex-Sum benchmark.

Instead of truncating long documents, LexLattice tiles the document's structure onto a 2D semantic lattice — rows are sections, columns are paragraphs — and runs a masked neural cellular automaton (NCA) over it, letting paragraph representations exchange information along both the reading order and the document hierarchy. A residual readout then scores every paragraph's salience, and a Plackett–Luce selection extracts the summary under a word budget.

frozen mT5-base encoder → paragraph states [N, 768]
    → tile onto (section × paragraph) lattice [768, 48, 32]  (masked; padding never updates)
    → + language-embedding channel
    → masked depthwise-conv 2D-NCA, 8 steps
    → un-tile → residual readout → salience score per paragraph
    → greedy Plackett–Luce selection up to the word budget

The frozen encoder is included in this repo (encoder/ subfolder) and is loaded automatically by the provided code. It is an mT5-base-architecture encoder whose transformer blocks carry the google/mt5-base pretrained weights, but whose input-embedding table was (deterministically) re-initialized in the training environment and therefore differs from stock mT5 — see Encoder provenance below. Only the consolidator (NCA + language embedding + readout) is trained; the encoder is frozen.

⚠️ Do not substitute google/mt5-base for the bundled encoder: the consolidator checkpoints were trained on the bundled weights and are not compatible with stock mT5 embeddings.

Checkpoints

File	Training
consolidator_2d.safetensors	Supervised warm-start: Plackett–Luce NLL on greedy-oracle extraction orderings (+ BCE membership auxiliary), trained on all 24 languages
consolidator_2d_rl.safetensors	The above, fine-tuned with RLOO (REINFORCE Leave-One-Out) against a blended ROUGE-1/2/Lsum reward, z-normalized per language

Original PyTorch .pt state dicts are included for exact reproducibility of the paper pipeline.

Usage

# pip install torch transformers safetensors huggingface_hub sentencepiece
from huggingface_hub import hf_hub_download
import importlib.util

path = hf_hub_download("lalala512/lexlattice-2d-nca", "modeling_lexlattice.py")
spec = importlib.util.spec_from_file_location("modeling_lexlattice", path)
mod = importlib.util.module_from_spec(spec); spec.loader.exec_module(mod)

model = mod.LexLattice.from_pretrained(
    "lalala512/lexlattice-2d-nca",
    checkpoint="consolidator_2d_rl.safetensors",   # or consolidator_2d.safetensors
)

# A document is a list of sections; each section is a list of paragraph strings.
sections = [
    ["Paragraph 1 of the preamble...", "Paragraph 2 of the preamble..."],
    ["Article 1 text...", "Article 2 text..."],
]
summary = model.summarize(sections, lang="english", budget_words=600)

# Or get raw per-paragraph salience scores:
texts, scores = model.score_document(sections, lang="english")

Splitting a raw document into sections/paragraphs is up to you; the paper uses a rule-based multilingual parser for EUR-Lex acts (see the project notebooks). Any reasonable hierarchical segmentation works — the lattice clamps overflow beyond 48 sections × 32 paragraphs into the boundary cells (clamp_merge).

Results

ROUGE F1 (×100), macro-averaged over the 24 EUR-Lex-Sum test languages, extractive selection at the reference word budget:

Model	R-1	R-2	R-Lsum
Random	45.22	16.34	39.39
Lead	49.85	21.10	43.60
BM25-centroid	41.46	15.93	34.44
1D-NCA (reading order only, ablation)	48.98	21.87	42.86
2D-NCA (this repo, supervised)	50.60	23.27	44.31
2D-NCA + RLOO (this repo)	50.67	22.92	44.33

The 2D lattice beats the otherwise-identical 1D-NCA ablation on every metric, i.e. the section axis carries signal beyond reading order. Per-language tables and further ablations (language embedding, cross-lingual zero-shot transfer, parse-noise sensitivity, LLM-judge) are in the accompanying paper.

Model details

Parameters	1,797,377 (consolidator) + frozen ~390M encoder (bundled in encoder/)
Encoder	mT5-base architecture (bundled; see Encoder provenance), masked mean pooling, ≤256 tokens per paragraph
Lattice	48 × 32 (section × paragraph), overflow clamp_merge
NCA	depthwise 3×3 perception + 1×1 MLP (hidden 256), 8 steps, deterministic
Readout	Linear(1536→768) → GELU → Linear(768→1) on [original ⊕ consolidated] states
Languages	24 official EU languages, learned language embedding (24 × 768)
Training data	EUR-Lex-Sum train splits, all 24 languages

Encoder provenance

The training environment loaded google/mt5-base through a code path that left the encoder's input-embedding table re-initialized (fixed seed) instead of inheriting mT5's pretrained embeddings, while all transformer blocks loaded correctly. Because the initialization was deterministic, the exact encoder used throughout training and evaluation has been recovered and is shipped in encoder/ — the checkpoints in this repo are bit-exact with the reported results when used with it. Practical consequences:

• All reported numbers were obtained with this encoder; they are internally consistent and reproducible from this repo alone.
• The paragraph representation is a contextual transform of a shared random projection of the mT5 SentencePiece vocabulary, not of mT5's pretrained embedding semantics. Results with a correctly-loaded mT5 embedding table may differ (plausibly improve) and will be explored in a future revision.

Intended use & limitations

• Extractive summarization of hierarchically structured legal/regulatory documents; scores and selects existing paragraphs, never generates text — so it cannot paraphrase or compress within a paragraph.
• Trained on EUR-Lex acts; transfer to other document genres is untested.
• The language embedding requires one of the 24 training languages. For unseen languages, see the cross-lingual transfer checkpoints in the paper (trained with the language embedding off).
• Not legal advice; summaries may omit legally material provisions. Always consult the full act.

Citation

Anonymous submission under review at ACL — citation information withheld during the review period.