RoBERTa Model - Custom Implementation

Overview

Model Architecture

Core Components

1. Embeddings Layer

• Token Embeddings: Maps input tokens to dense vectors (VOCAB_SIZE × HIDDEN_SIZE)
• Positional Embeddings: Encodes absolute position information (MAX_LEN × HIDDEN_SIZE)
• Layer Normalization: Normalizes embedding outputs for stable training
• Dropout: Regularization with rate = 0.1

2. Encoder Blocks (Stacked Transformer Encoders)

• Multi-Head Self-Attention:
  • Number of heads: 8
  • Allows the model to attend to different parts of the sequence
  • Uses key padding mask to ignore padding tokens
• Feed-Forward Network (FFN):
  • Linear layer: HIDDEN_SIZE → FFN_DIM (256 → 1024)
  • Activation: GELU
  • Linear layer: FFN_DIM → HIDDEN_SIZE (1024 → 256)
  • Expands and compresses information
• Residual Connections & Layer Normalization:
  • Applied after attention and FFN layers
  • Stabilizes deep network training

3. MLM (Masked Language Model) Head

• Dense Layer: HIDDEN_SIZE → HIDDEN_SIZE (256 → 256)
• GELU Activation: Non-linear transformation
• Layer Normalization: Output normalization
• Decoder: HIDDEN_SIZE → VOCAB_SIZE (256 → 16000)
• Outputs logits for vocabulary prediction

Architecture Summary

Input IDs (batch_size, seq_len)
    ↓
Embeddings Layer (Token + Position + LayerNorm + Dropout)
    ↓
Encoder Block 1 (Multi-Head Attention + FFN + Residual + LayerNorm)
    ↓
Encoder Block 2 (Multi-Head Attention + FFN + Residual + LayerNorm)
    ↓
Encoder Block 3 (Multi-Head Attention + FFN + Residual + LayerNorm)
    ↓
Encoder Block 4 (Multi-Head Attention + FFN + Residual + LayerNorm)
    ↓
MLM Head (Dense + GELU + LayerNorm + Decoder)
    ↓
Output Logits (batch_size, seq_len, vocab_size)

Model Hyperparameters

Data

Data Sources

1. Wikipedia 2023
   • Source: wikimedia/wikipedia (20231101.en split)
   • Samples used: 100,000
   • High-quality encyclopedic text
2. Alpaca Dataset
   • Source: tatsu-lab/alpaca
   • Samples used: 100,000
   • Instruction-following examples and outputs
3. TinyStories
   • Source: roneneldan/TinyStories
   • Samples used: 100,000
   • Narrative and story-based text

Preprocessing Pipeline

1. Corpus Generation

• Combined datasets are interleaved to create a unified corpus
• Text from all sources is concatenated into corpus.txt
• One document per line

2. Tokenization

• Tokenizer: SentencePiece (BPE model type)
• Vocabulary size: 16,000 tokens
• Special tokens:
  • <pad> (ID: 0) - Padding token
  • <unk> (ID: 1) - Unknown token
  • BOS (ID: 2) - Beginning of sequence
  • EOS (ID: 3) - End of sequence
  • [MASK] (custom) - Masking token

3. Chunking Strategy

• Sequence length: 128 tokens (including BOS and EOS)
• Content length: 126 tokens (128 - 2 for BOS/EOS)
• Each text is split into overlapping chunks of 126 tokens
• Chunks are padded to length 128 with pad tokens
• Saved as PyTorch tensors in chunks/ directory (10,000 chunks per file)

Dataset Statistics

• Total chunks generated: 200+ files
• Samples per file: 10,000
• Total training samples: 2,000,000+
• Sequence length: 128 tokens

Training

Masked Language Modeling (MLM) Objective

MLM Process (15% masking probability):

1. Token Selection: 15% of tokens are randomly selected for masking
2. Masking Strategy (applied to selected tokens):
   • 80% replaced with [MASK] token
   • 10% replaced with a random token from vocabulary
   • 10% left unchanged
3. Loss Calculation:
   • Only masked positions contribute to the loss
   • Non-masked positions have labels set to -100 (ignored)
   • Cross-entropy loss is computed on masked token predictions

Example:

Original:     "the quick brown fox jumps over the lazy dog"
Tokens:       [2, 101, 102, 103, 104, 105, 106, 107, 108, 109, 3, 0, ...]
Masked:       [2, [MASK], 102, [MASK], 104, random_token, 106, 107, [MASK], 109, 3, 0, ...]
Labels:       [-100, 101, -100, 104, -100, 105, -100, -100, -100, -100, -100, -100, ...]
                      ↑              ↑              ↑              ↑
                   Predict original tokens at these positions

Training Configuration

• Loss function: CrossEntropyLoss
• Optimizer: AdamW
• Batch size: 4
• Dataset split: 200,000 training samples (subset for testing)

Project Structure

d:\BERT/
├── main.ipynb                      # Main training and evaluation notebook
├── requirements.txt                # Python dependencies
├── README.md                       # This file
│
├── corpus.txt                      # Combined text corpus
├── chunks.txt                      # Chunk metadata
├── english_tokenizer.model         # SentencePiece tokenizer model
├── english_tokenizer.vocab         # Tokenizer vocabulary file
│
├── chunks/                         # Preprocessed token chunks
│   ├── all_chunks.pt              # All chunks combined
│   ├── chunks_0.pt
│   ├── chunks_1.pt
│   ├── ... (up to chunks_200+.pt)
│
├── models/                         # Trained model checkpoints
│   └── (saved model files)
│
├── model_checkpoints/              # Training checkpoints
│   └── (checkpoint files)
│
├── sentence_embeddings.html        # 3D visualization of sentence embeddings
├── word_embeddings.html            # 3D visualization of word embeddings
├── sentence_embeddings_2d.png      # 2D PCA visualization
│
└── bert_venv/                      # Python virtual environment

Usage

Dependencies

torch                    # Deep learning framework
transformers            # For tokenizer utilities
datasets                # For loading public datasets
sentencepiece           # Tokenization
huggingface_hub         # Model hub integration
matplotlib              # Visualization
scikit-learn            # PCA visualization
plotly                  # Interactive plots
pandas                  # Data manipulation

Installation

Training

1. Load and authenticate with Hugging Face (requires HF_TOKEN in .env)
2. Generate corpus from combined datasets
3. Train SentencePiece tokenizer
4. Create text chunks and save as tensors
5. Initialize RoBERTa model
6. Train with masked language modeling objective
7. Save checkpoints and evaluate

Evaluation

• Mask accuracy: Percentage of correctly predicted masked tokens
• Embedding analysis:
  • Word embeddings visualization (3D and 2D PCA)
  • Sentence embeddings visualization
  • Similarity matrices between sequences

Key Features

Implementation Details

Special Design Choices

1. Dynamic masking guarantee: If a sequence has no masked tokens, at least one token is randomly selected for masking to ensure effective learning
2. Attention masking: Padding tokens are properly masked to prevent attention to pad positions
3. Residual connections: Applied in both attention and FFN blocks for training stability
4. GELU activation: Used instead of ReLU for smoother gradients and better performance
5. Layer normalization placement: Applied after each sublayer (post-norm) for numerical stability

Performance Metrics

• Training loss: Cross-entropy loss on masked tokens
• Masking accuracy: Correct predictions / total masked tokens
• Per-layer hidden states: For analysis and visualization

References

• Original RoBERTa Paper: RoBERTa: A Robustly Optimized BERT Pretraining Approach
• BERT: BERT: Pre-training of Deep Bidirectional Transformers
• Attention Is All You Need: Transformer Architecture

License

Author