docs/architecture.md

LexiMind Architecture

Overview

LexiMind couples a from-scratch Transformer implementation with a modern data and inference stack. The project consists of three major layers:

1. Data & Tokenization – HuggingFace tokenizer wrapper with tensor-aware batching and T5-specific decoder input preparation.
2. Model Composition – the bespoke encoder/decoder stack with task heads assembled via MultiTaskModel, plus models.factory.build_multitask_model to rebuild the network from configuration files.
3. Inference & Serving – a multi-task pipeline capable of summarization, emotion, and topic classification; surfaced through a CLI and Gradio UI.

Custom Transformer Stack

The custom Transformer is designed with modern architectural choices while maintaining compatibility with pre-trained weights from Google's FLAN-T5.

Architecture Highlights

• Pre-Layer Normalization (Pre-LN): RMSNorm applied before each sublayer for stable training
• RMSNorm: More efficient than LayerNorm (no mean computation, no bias parameters)
• FlashAttention: Via PyTorch 2.0's F.scaled_dot_product_attention for O(N) memory
• Learned Positional Embeddings: Trainable position representations (randomly initialized)
• Multi-Head Attention: 12 heads with optional LoRA adapters and RoPE support

Weight Loading from FLAN-T5

The factory.py module loads weights from FLAN-T5-base, which uses a compatible Pre-LN architecture:

• Token embeddings: Shared between encoder and decoder
• Attention projections: Q, K, V, O weights (bias initialized to zero since T5 has no attention bias)
• FFN weights: wi_1 → linear1, wo → linear2 (T5 uses gated FFN; we use the up/down projections)
• RMSNorm weights: Direct transfer (both use RMSNorm without bias)
• LM head: Loaded from T5's lm_head

Note: T5 uses relative position bias computed in attention, not absolute embeddings. Our learned positional embeddings are randomly initialized and train quickly during fine-tuning.

File Structure

• src/models/encoder.py – TransformerEncoder with Pre-LN RMSNorm blocks
• src/models/decoder.py – TransformerDecoder with KV-cache for efficient generation
• src/models/attention.py – Multi-Head Attention with FlashAttention, LoRA, and RoPE support
• src/models/heads.py – ClassificationHead (mean pooling) and LMHead (with weight tying)
• src/models/multitask.py – Routes inputs to task-specific heads
• src/models/factory.py – Builds models and loads FLAN-T5 weights

Data, Tokenization, and Datasets

• src/data/tokenization.py wraps AutoTokenizer (configured for FLAN-T5) to provide tensor-aware batching and helper utilities for decoder input shifting.
• src/data/dataset.py and src/data/dataloader.py define strongly typed dataset containers and task-specific collators.
• scripts/download_data.py fetches and processes training data from HuggingFace datasets.

Training Datasets

Task	Dataset	Size	Labels
Summarization	BookSum + arXiv	~90K	Text→Summary
Emotion	GoEmotions	~43K	28 emotions (multi-label)
Topic	Books + Papers	~50K	8 categories (Fiction, Science, Technology, etc.)
Books	Gutenberg (prose chunks)	~30K	Literary text

T5 Tokenizer Differences

• Vocab size: 32,128 tokens (SentencePiece)
• Special tokens: pad=0, eos=1 (no explicit BOS; decoder starts with pad token)
• Subword tokenization: Unigram-based (vs BART's BPE)

Training Pipeline

• src/training/trainer.py coordinates multi-task optimization with:
  • Mixed precision training (bfloat16 on Ampere/Ada GPUs)
  • Gradient accumulation for larger effective batch sizes
  • Per-task loss weighting and label smoothing
  • Early stopping based on validation loss
  • Cosine learning rate schedule with warmup
• torch.compile: JIT compilation with Inductor backend for 20-40% speedup
• Metrics in src/training/metrics.py include accuracy, multi-label F1, and ROUGE-like overlap

Inference & Serving

• src/inference/pipeline.py exposes summarization, emotion, and topic predictions with shared pre-processing, generation, and thresholding logic.
• src/inference/factory.py rebuilds the full pipeline using the exported tokenizer artifact
• The CLI (scripts/inference.py) drives the pipeline from the command line
• Gradio demo (scripts/demo_gradio.py) provides an interactive web interface

Key Decisions

• Custom Transformer + Pre-trained Weights: Building from scratch demonstrates deep understanding while leveraging FLAN-T5's language knowledge
• Pre-LN RMSNorm: Modern architecture used by LLaMA, T5 v1.1, and other 2023-2025 models
• Simplified Training: Removed NaN detection and gradient monitoring (Windows workarounds no longer needed on WSL/Linux)
• Clean Dataset Pipeline: AG News (4 clean categories) instead of Yahoo Answers (10 messy categories); BookSum for literary summarization
• Tokenizer Artifact Preference: Inference favors artifacts/hf_tokenizer for reproducibility