💬 NLP

DEPTH: Discourse Education through Pre-Training Hierarchically

The paper introduces DEPTH, a hierarchical encoder-decoder model that enhances language models' discourse understanding through a novel pre-training objective combining sentence un-shuffling and span-corruption, resulting in faster learning and superior performance on both discourse-specific and general natural language understanding benchmarks compared to standard T5.

Original authors: Zachary Bamberger, Ofek Glick, Chaim Baskin, Yonatan Belinkov

Published 2026-02-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zachary Bamberger, Ofek Glick, Chaim Baskin, Yonatan Belinkov

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to write stories. You give it a massive library of books (the internet) and tell it, "Read all of this and learn how to predict the next word."

Most modern language models (like the famous T5) are incredibly good at this. They learn grammar, vocabulary, and facts very quickly. However, they often struggle with the big picture. They might write a sentence that is grammatically perfect but makes no sense in the context of the paragraph, or they might get the order of events in a story completely wrong. They are like a student who memorized every word in a dictionary but doesn't understand how to write a coherent essay.

This paper introduces DEPTH, a new way of training these robots to understand the "flow" of a story, not just the words.

Here is how DEPTH works, explained through simple analogies:

1. The Problem: The "Word-By-Word" Robot

Imagine a robot trying to assemble a puzzle.

Standard Models (T5): They look at the puzzle piece by piece. They know what a "sky" piece looks like and what a "grass" piece looks like. But if you hand them the pieces in a random order, they might get confused about how the sky connects to the grass. They focus on the texture of the pieces (sub-word level) but miss the picture (sentence level).
The Goal: We want a robot that understands that "The sun rose" must come before "The birds started singing," even if the words are jumbled.

2. The Solution: The "Sentence-First" Approach (DEPTH)

The authors created a new training method called DEPTH. Think of it as giving the robot a new set of training exercises that force it to think in "chunks" (sentences) rather than just "bricks" (words).

They use two main tricks:

Trick A: The Sentence Shuffle (The "Scrambled Puzzle" Game)

Imagine you take a paragraph from a book, cut it into individual sentences, and throw them into a blender.

The Task: The robot has to look at the scrambled sentences and figure out the correct order to put them back together to make a story that makes sense.
Why it helps: This forces the robot to learn the relationships between sentences. It learns that "Once upon a time" usually starts a story, and "The end" usually finishes it. It stops just guessing the next word and starts understanding the narrative flow.

Trick B: The "Hidden Clue" Game (Span Corruption)

This is a standard trick used in AI training, but DEPTH uses it differently.

The Task: You take a sentence, hide a few words (like "The cat sat on the [MASK]"), and ask the robot to guess the missing word.
The Twist: In DEPTH, the robot has to do this while it is also trying to solve the scrambled sentence puzzle. It has to wear two hats at once: understanding the local details (the missing word) and the global structure (the sentence order).

3. The Secret Sauce: The "Sentence Token"

To make this work, the authors gave the robot special "ID cards" for sentences.

Imagine every sentence in a document gets a colored tag (like a red tag for Sentence 1, a blue tag for Sentence 2).
The robot is trained to pay attention to these tags. It learns that all the words inside the "Red Tag" belong together, and it needs to figure out how the "Red Tag" connects to the "Blue Tag."
This creates a hierarchy: The robot understands the small details (words) inside the big details (sentences).

4. The Results: A Faster, Smarter Learner

The researchers tested DEPTH against the standard T5 model.

Speed: DEPTH learned much faster. It was like a student who, after reading just a few chapters, could already write a coherent story, while the standard student was still struggling with sentence structure.
Understanding: When asked to fix jumbled stories or check if a paragraph made sense, DEPTH was much better. It understood the "story logic" much more deeply.
No Trade-offs: Usually, when you teach a robot a new skill, it forgets an old one. But DEPTH got better at understanding stories without getting worse at grammar or vocabulary.

The Big Picture Analogy

Think of learning a language like learning to drive a car.

Standard Models are great at knowing what the pedals, steering wheel, and mirrors do (the mechanics/words).
DEPTH teaches the driver how to navigate traffic, understand road signs, and plan a route from Point A to Point B (the discourse/flow).

By teaching the robot to "drive" the whole sentence structure, not just the individual words, DEPTH creates a language model that feels more human because it understands how ideas connect to form a complete, logical thought.

In short: DEPTH is a training method that stops AI from just memorizing words and starts teaching it how to tell a good story.

1. Problem Statement

Large Language Models (LMs), particularly encoder-decoder architectures like T5, often struggle with discourse-level understanding. While pre-training data contains abundant discourse patterns (coherence, cohesion, narrative flow), standard pre-training objectives (like span corruption) focus primarily on sub-word and sentence-level semantics.

Limitations of Previous Approaches: Early attempts to inject discourse awareness, such as Next Sentence Prediction (NSP) in BERT or Sentence Order Prediction (SOP) in ALBERT, proved too simplistic and hindered effective learning. Conversely, decoder-only models (GPT series) excel at generation but often fail at discourse-oriented benchmarks without costly alignment phases.
The Gap: There is a need for a pre-training objective that explicitly teaches encoder-decoder models to understand the hierarchical relationships between sentences and the global structure of a document, without sacrificing their generative or general NLU capabilities.

2. Methodology: DEPTH

The authors propose DEPTH (Discourse Education through Pre-Training Hierarchically), a hierarchical encoder-decoder model that integrates discourse-aware objectives directly into the pre-training phase.

Core Components

Hierarchical Tokenization:
- DEPTH introduces a new tokenizer that segments input text into sentences using NLTK.
- It augments the vocabulary with special sentence tokens: <SENT_i> (indicating sentence index $i$ ) and <EOSEN> (End of Sentence).
- Unlike the SLM (Sentence-level Language Model) which used an auxiliary pointer-decoder, DEPTH is a full encoder-decoder that predicts these sentence tokens directly.
Dual Pre-Training Objectives:
The model is trained simultaneously on two tasks:
- Span Corruption (Reconstruction): Similar to T5, random spans of tokens are masked and replaced with sentinel tokens. The model must reconstruct the original text.
- Sentence Un-Shuffling: The order of sentences in the input is randomly shuffled (with probability $p=0.5$ ). The model must reconstruct the original sentence order.
- Note: To prevent the sentinel tokens from giving away sentence positions (a flaw in standard T5 span masking), DEPTH randomly samples sentinel token IDs rather than using a fixed decreasing sequence.
Hierarchical Attention Mechanism:
- Encoder Self-Attention: Non-sentence tokens can attend to all tokens. However, <SENT_i> tokens can only attend to tokens within their own sentence. This forces the model to learn a robust representation of individual sentences before integrating them.
- Decoder Self-Attention: <SENT_i> tokens can only attend to past sentence tokens, enforcing a causal flow at the sentence level.
- Cross-Attention: Non-sentence tokens in the decoder attend to the full encoder input, while sentence tokens in the decoder can only attend to sentence tokens in the encoder. This encourages the model to use sentence-level representations to predict the next sentence.
Loss Formulation:
The total loss is a token-averaged cross-entropy combining both the reconstruction loss (span masking) and the sentence un-shuffling loss.

3. Experimental Setup

Baselines: The authors compare DEPTH against T5-Base (220M parameters).
Datasets: Pre-training is conducted on the C4 dataset.
Training Regimes:
1. From Scratch (FS): Both models are randomly initialized.
2. Continuous Pre-Training (CPT): Both models start from a pre-trained T5-Base checkpoint and continue training on C4.
Evaluation Benchmarks:
- GLUE: General NLU tasks (CoLA, SST-2, MNLI).
- DiscoEval: Specific discourse tasks (Sentence Permutation - SP, Discourse Coherence - DC).
- Natural Instructions (NI): Instruction-following capabilities.
Constraints: Due to limited compute, the authors did not use "example packing" (concatenating multiple examples into a single batch), which is standard for T5 but can introduce noise. They used AdamW optimizer and specific learning rate schedules.

4. Key Results

Pre-Training Efficiency

Faster Convergence: DEPTH achieves lower validation loss than T5 during pre-training, even when balancing the additional sentence un-shuffling objective.
Reconstruction Performance: When isolating the span-corruption loss, DEPTH still outperforms T5, suggesting the hierarchical structure aids general representation learning.

Downstream Performance

Discourse Tasks (DiscoEval):
- From Scratch: DEPTH significantly outperforms T5 on Discourse Coherence (DC) tasks, demonstrating robustness in understanding narrative structures.
- Sentence Permutation (SP): DEPTH struggles more than T5 on SP tasks in the FS setting (accuracy $\le$ 5% on shuffled tokens), indicating that sentence un-shuffling is a harder task than previous discourse objectives (like NSP). However, DEPTH makes more "reasonable" errors (coherent but incorrect orderings) compared to T5.
General NLU (GLUE):
- In the FS setting, DEPTH learns GLUE tasks faster than T5.
- In the CPT setting, DEPTH performs comparably to T5, showing that adding discourse objectives does not degrade existing NLU capabilities.
Instruction Following (NI):
- FS: DEPTH shows a massive performance leap between 16k and 32k steps, outperforming T5 significantly.
- CPT: DEPTH underperforms T5, likely due to a domain shift from the original T5 pre-training to the shuffled-sentence objective.

5. Key Contributions

Novel Pre-Training Objective: Introduces a unified encoder-decoder framework that combines span corruption with sentence un-shuffling, moving beyond simple next-sentence prediction.
Hierarchical Architecture: Proposes a specific attention masking scheme that enforces sentence-level isolation in the encoder and causal sentence flow in the decoder, allowing the model to learn latent sentence representations.
Efficiency: Demonstrates that discourse-aware pre-training allows models to learn semantic and discourse representations faster than standard T5, requiring fewer training steps to reach high performance on downstream tasks.
Preservation of Capabilities: Shows that enhancing discourse capabilities does not come at the cost of general NLU performance.

6. Significance and Future Work

Significance: The paper challenges the notion that decoder-only models are the only path to strong generative capabilities. It proves that encoder-decoder models, when pre-trained with hierarchical discourse objectives, can achieve superior discourse understanding and faster learning dynamics.
Limitations: The study was constrained by computational resources (smaller batch sizes, no example packing, fewer total tokens than T5's original training), which may explain why DEPTH didn't fully match the absolute performance of the original T5 on all metrics.
Future Directions:
- Applying DEPTH to Retrieval-Augmented Generation (RAG) using sentence "chunks."
- Extending the hierarchy to higher-level units (paragraphs, chapters).
- Scaling to larger model sizes to test if discourse capabilities amplify with capacity.
- Investigating alternative sentence-level tasks (e.g., Next Sentence Prediction) as replacements for un-shuffling.

In conclusion, DEPTH establishes that explicitly modeling discourse structure during pre-training is a highly effective strategy for improving the coherence and narrative understanding of language models, offering a promising path for more human-like language comprehension.