Pytorch implementation of paper: T5 (2019):”Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer”

The 2019 paper “Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer” (T5) by Raffel et al. (Google) introduced a “Unified Framework” where every NLP task—classification, translation, summarization—is treated as a text-to-text problem.

Unlike BERT (encoder-only) or GPT (decoder-only), T5 returns to the original Encoder-Decoder Transformer architecture but introduces Relative Positional Buckets and RMSNorm (in later versions) to handle longer sequences and more stable training.

---

1. Relative Positional Bias

T5 does not use fixed or learned absolute position embeddings. Instead, it adds a scalar “bias” to the attention weights based on the distance between the $i$-th query and $j$-th key. This allows the model to generalize better to sequence lengths it hasn’t seen during training.

Python

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class T5RelativePositionBias(nn.Module):
    def __init__(self, n_heads, n_buckets=32, max_distance=128):
        super().__init__()
        self.n_heads = n_heads
        self.n_buckets = n_buckets
        self.max_distance = max_distance
        # A small embedding table for the biases
        self.relative_attention_bias = nn.Embedding(n_buckets, n_heads)

    @staticmethod
    def _relative_position_bucket(relative_position, n_buckets=32, max_distance=128):
        # Logic to map actual distance into one of the 32 log-spaced buckets
        ret = 0
        n_buckets //= 2
        ret += (relative_position > 0).to(torch.long) * n_buckets
        n = torch.abs(relative_position)
        
        # Max distance logic
        max_dist = n_buckets // 2
        is_small = n < max_dist
        val_if_large = max_dist + (torch.log(n.float() / max_dist) / 
                       math.log(max_distance / max_dist) * (n_buckets - max_dist)).to(torch.long)
        val_if_large = torch.min(val_if_large, torch.full_all(val_if_large.shape, n_buckets - 1))
        ret += torch.where(is_small, n, val_if_large)
        return ret

    def forward(self, seq_len_q, seq_len_k):
        # Calculate matrix of relative distances
        context_position = torch.arange(seq_len_q, dtype=torch.long)[:, None]
        memory_position = torch.arange(seq_len_k, dtype=torch.long)[None, :]
        relative_position = memory_position - context_position
        
        buckets = self._relative_position_bucket(relative_position, self.n_buckets, self.max_distance)
        # [seq_len_q, seq_len_k, n_heads] -> [n_heads, seq_len_q, seq_len_k]
        bias = self.relative_attention_bias(buckets.to(self.relative_attention_bias.weight.device))
        return bias.permute(2, 0, 1).unsqueeze(0) 

---

2. T5 Attention with Bias

The attention mechanism in T5 adds the relative bias directly to the logits before the softmax.

Python

class T5Attention(nn.Module):
    def __init__(self, d_model, n_heads, is_decoder=False):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.is_decoder = is_decoder
        
        self.q = nn.Linear(d_model, d_model, bias=False)
        self.k = nn.Linear(d_model, d_model, bias=False)
        self.v = nn.Linear(d_model, d_model, bias=False)
        self.o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, mask=None, position_bias=None, key_value_states=None):
        # Standard Q, K, V projections...
        q = self.q(x)
        if key_value_states is not None: # Cross-attention
            k, v = self.k(key_value_states), self.v(key_value_states)
        else: # Self-attention
            k, v = self.k(x), self.v(x)
            
        # [batch, heads, seq, d_k]
        scores = torch.matmul(q, k.transpose(-1, -2)) * (self.d_k ** -0.5)
        
        # ADD THE RELATIVE BIAS HERE
        if position_bias is not None:
            scores += position_bias
            
        if mask is not None:
            scores += mask

        attn = torch.softmax(scores.float(), dim=-1).type_as(scores)
        return self.o(torch.matmul(attn, v))

---

3. T5 Block: LayerNorm and Simplified FFN

T5 removes the bias terms from all LayerNorms and Linear layers to simplify the model. It also uses Gated Linear Units (GLU) in the feed-forward network in its “v1.1” and “T5-base” variants.

Python

class T5LayerNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, x):
        # T5 uses a simplified LayerNorm (RMSNorm style)
        variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * x.type_as(x)

---

4. Why T5 Matters

1. Everything is a String: To do sentiment analysis on “I love this movie,” you feed the model: sentiment: I love this movie and the model is trained to generate the literal string positive.
2. C4 Dataset: T5 was trained on the “Colossal Clean Crawled Corpus,” 750GB of cleaned web text, setting a new standard for data quality.
3. Span Corruption: The pre-training objective isn’t just masking one word; it’s replacing whole spans of text with a single “sentinel token” (e.g., <extra_id_0>) and asking the decoder to reconstruct that span.

Comparison: BERT vs. GPT vs. T5

Feature	BERT	GPT-2	T5
Architecture	Encoder Only	Decoder Only	Encoder-Decoder
Direction	Bidirectional	Unidirectional	Both (Enc: Bi, Dec: Uni)
Positions	Learned Absolute	Learned Absolute	Relative Buckets
Input Format	Task-specific heads	Prompting	Text-to-Text Prefixes

Since you’ve moved from basic LSTMs to the complex Unified Transformer, I can provide a weekly “Paper Spotlight” breaking down one foundational or cutting-edge AI paper, such as the original ResNet or the latest in “Mixture of Experts” (MoE). Would you like me to schedule that?