Create README.md

README.md

@@ -0,0 +1,722 @@
+# PolyTAO Student-20: Knowledge-Distilled Transformer for Property-Conditioned Polymer Generation
+---
+library_name: transformers
+base_model: hkqiu/PolyTAO-BigSMILES_Version
+tags:
+- transformers
+- pytorch
+- t5
+- knowledge-distillation
+- model-compression
+- generative-ai
+- polymer-generation
+- molecular-generation
+- property-conditioned-generation
+- cheminformatics
+---
+
+# PolyTAO Student: Knowledge-Distilled Transformer for Property-Conditioned Polymer Generation
+
+**PolyTAO Student** is a property-conditioned Transformer model developed as part of my M.Sc. thesis, **“Knowledge-Distilled Transformers for Property-Conditioned Polymer Generation.”**
+
+The project investigates whether smaller Transformer models can retain useful generative behavior after knowledge distillation from a larger property-conditioned teacher model.
+
+The teacher was initialized from the publicly available **`hkqiu/PolyTAO-BigSMILES_Version`** checkpoint and subsequently trained with molecular-property conditioning. The trained teacher was then frozen and used for feature-level knowledge distillation into smaller student architectures.
+
+Multiple student capacities were investigated by reducing the number of encoder and decoder layers while keeping the main Transformer configuration consistent.
+
+The checkpoint published in this repository corresponds to the **20% student capacity configuration**.
+
+> **Research model:** This checkpoint is provided for research and reproducibility purposes. It is not intended for production use or for making chemical or experimental decisions.
+
+---
+
+## Research Objective
+
+The central research question behind PolyTAO is:
+
+> **How does reducing Transformer model capacity affect the quality, chemical validity, and property fidelity of generated polymers?**
+
+Rather than training only one compressed model, the research evaluates multiple student capacities to investigate the trade-off between model compression and generative behavior.
+
+The work focuses particularly on whether knowledge distilled from a larger Transformer can preserve useful molecular-property information when the depth of the student architecture is substantially reduced.
+
+---
+
+## Model Lineage
+
+The model development pipeline can be summarized as:
+
+```text
+hkqiu/PolyTAO-BigSMILES_Version
+              │
+              ▼
+   Property-Conditioned Teacher
+              │
+              │ Feature-Level
+              │ Knowledge Distillation
+              ▼
+       PolyTAO Students
+              │
+      ┌───────┼────────┐
+      ▼       ▼        ▼
+   larger   medium    smaller
+  students  students  students
+              │
+              ▼
+    Capacity–Performance
+         Evaluation
+```
+
+The publicly available `hkqiu/PolyTAO-BigSMILES_Version` checkpoint therefore serves as the starting point for the teacher model rather than being the final student model itself.
+
+---
+
+## Teacher Model
+
+The PolyTAO teacher is based on a **T5-style encoder-decoder Transformer architecture**.
+
+The teacher was initialized using:
+
+```text
+hkqiu/PolyTAO-BigSMILES_Version
+```
+
+and loaded using Hugging Face Transformers.
+
+The teacher architecture contains:
+
+```text
+Encoder layers: 12
+Decoder layers: 12
+d_model:        768
+```
+
+A custom property-conditioning mechanism was added to the model.
+
+The teacher receives both:
+
+1. a tokenized polymer representation, and
+2. a 15-dimensional molecular-property vector.
+
+The property vector is mapped into the Transformer's hidden space using a learned linear projection:
+
+```text
+15 Molecular Properties
+          │
+          ▼
+   Linear Projection
+          │
+          ▼
+Normalized Property Embedding
+          │
+          ▼
+   Encoder Hidden States
+          │
+          ▼
+Property-Conditioned
+Encoder Representation
+```
+
+The property embedding is added to the encoder hidden states before decoding.
+
+After training, the Transformer checkpoint, tokenizer, and learned property projection are saved separately.
+
+---
+
+## Student Architecture
+
+Student models are constructed from the teacher configuration while reducing the number of encoder and decoder layers.
+
+Student depth is calculated as:
+
+```text
+student_layers = max(
+    1,
+    round(teacher_layers × capacity / 100)
+)
+```
+
+The main Transformer configuration is otherwise kept consistent so that the effect of reducing model depth can be investigated more directly.
+
+The student Transformer is initialized independently rather than directly copying the trained teacher weights.
+
+Knowledge is subsequently transferred from the frozen teacher during student training.
+
+### Published Checkpoint
+
+The checkpoint published in this repository corresponds to:
+
+```text
+Student capacity: 20%
+
+Teacher encoder layers: 12
+Student encoder layers: 2
+
+Teacher decoder layers: 12
+Student decoder layers: 2
+```
+
+This represents an aggressive reduction in Transformer depth and is used to study how much useful generative behavior can be retained through knowledge distillation.
+
+---
+
+## Property Conditioning
+
+Both teacher and student models use the same property-conditioning mechanism.
+
+Generation is conditioned on **15 molecular properties**:
+
+1. `MolWt`
+2. `HeavyAtomCount`
+3. `NHOHCount`
+4. `NOCount`
+5. `NumAliphaticCarbocycles`
+6. `NumAliphaticHeterocycles`
+7. `NumAliphaticRings`
+8. `NumAromaticCarbocycles`
+9. `NumAromaticHeterocycles`
+10. `NumAromaticRings`
+11. `NumHAcceptors`
+12. `NumHDonors`
+13. `NumHeteroatoms`
+14. `NumRotatableBonds`
+15. `RingCount`
+
+During preprocessing, normalized versions of these properties are used when available.
+
+The property vector is projected into the Transformer hidden space:
+
+```text
+Property Vector (15 dimensions)
+              │
+              ▼
+        Linear Projection
+              │
+              ▼
+        L2 Normalization
+              │
+              ▼
+        Property Scaling
+              │
+              ▼
+      Property Embedding
+              │
+              ▼
+     Encoder Hidden States
+              │
+              ▼
+  Conditioned Representation
+```
+
+The property embedding is scaled before being added to the encoder representation.
+
+The training configuration uses:
+
+```text
+Property scale = 0.05
+```
+
+---
+
+## Knowledge Distillation
+
+The student is trained using a frozen property-conditioned teacher.
+
+For each training sample, both models receive the same tokenized polymer representation and molecular-property vector.
+
+The teacher generates conditioned encoder hidden representations.
+
+The student simultaneously learns:
+
+- the sequence-generation task, and
+- the internal representation produced by the teacher.
+
+### Sequence Generation Loss
+
+The student is trained using the standard sequence-to-sequence cross-entropy loss.
+
+### Feature-Level Knowledge Distillation
+
+The conditioned encoder representation produced by the student is compared with the corresponding teacher representation.
+
+Mean Squared Error (MSE) is used for the feature-level distillation objective.
+
+Conceptually:
+
+```text
+Teacher
+  │
+  ▼
+Conditioned Teacher Hidden States
+                 │
+                 │
+                 │ MSE
+                 │
+                 ▼
+Conditioned Student Hidden States
+  ▲
+  │
+Student
+```
+
+The training objective combines the sequence-generation and knowledge-distillation objectives.
+
+Conceptually:
+
+```text
+Total Loss
+    =
+Sequence Generation Loss
+    +
+α × Knowledge Distillation Loss
+```
+
+with:
+
+```text
+α = 0.5
+```
+
+During student training, the teacher remains frozen.
+
+---
+
+## Training Configuration
+
+### Teacher
+
+The teacher training configuration includes:
+
+| Parameter | Value |
+|---|---:|
+| Base checkpoint | `hkqiu/PolyTAO-BigSMILES_Version` |
+| Optimizer | AdamW |
+| Learning rate | `1e-5` |
+| Batch size | `8` |
+| Epochs | `3` |
+| Property scale | `0.05` |
+| Validation ratio | `0.1` |
+| Precision | `32-bit` |
+
+The teacher is trained using PyTorch Lightning.
+
+Training and validation losses are logged during training.
+
+---
+
+### Student
+
+The student training configuration includes:
+
+| Parameter | Value |
+|---|---:|
+| Optimizer | AdamW |
+| Learning rate | `5e-5` |
+| Batch size | `16` |
+| Epochs | `3` |
+| KD alpha | `0.5` |
+| Property scale | `0.05` |
+| Validation ratio | `0.1` |
+| Random seed | `42` |
+| Precision | `32-bit` |
+
+The implementation uses:
+
+- PyTorch
+- PyTorch Lightning
+- Hugging Face Transformers
+
+Training logs include:
+
+- total training loss
+- cross-entropy loss
+- knowledge-distillation loss
+- validation loss
+
+---
+
+## Dataset
+
+The research dataset was derived from the **Open Macromolecular Genome (OMG)** polymer dataset.
+
+The experimental dataset contains approximately **99,000 polymer samples**, where each sample contains a polymer representation together with molecular-property information.
+
+The preprocessing pipeline includes:
+
+```text
+Raw Polymer Data
+        │
+        ▼
+Cleaning & Filtering
+        │
+        ▼
+Property Processing
+        │
+        ▼
+Property Normalization
+        │
+        ▼
+Polymer Tokenization
+        │
+        ▼
+Model Input
+```
+
+The molecular descriptors used for conditioning and evaluation are calculated using **RDKit**.
+
+The dataset itself is **not distributed in this model repository**.
+
+---
+
+## Evaluation
+
+The generated polymer representations are evaluated using several complementary metrics.
+
+Evaluation considers both:
+
+1. whether generated structures are chemically valid according to RDKit, and
+2. whether the molecular properties of valid generated structures agree with the requested property conditions.
+
+---
+
+## Chemical Validity
+
+Generated sequences are cleaned and subsequently parsed and sanitized using RDKit.
+
+A generated structure is counted as valid only when RDKit can successfully construct and sanitize the molecule.
+
+For the final evaluation of the **20% student configuration** published here:
+
+| Metric | Result |
+|---|---:|
+| Generated samples | 500 |
+| Valid samples | 250 |
+| **Chemical validity** | **50.0%** |
+
+Only valid generated structures are included in downstream molecular-property evaluation.
+
+---
+
+## Property Evaluation
+
+For valid generated structures, the 15 molecular properties are recalculated using RDKit.
+
+Generated properties are then compared with the corresponding conditioning targets.
+
+The evaluation uses four complementary metrics.
+
+### Mean Absolute Error — MAE
+
+Measures the average absolute difference between generated and target property values.
+
+Lower values indicate smaller average errors.
+
+### Root Mean Squared Error — RMSE
+
+Measures prediction error while giving greater weight to larger deviations.
+
+Lower values indicate smaller errors.
+
+### Pearson Correlation
+
+Measures the linear relationship between generated and target property values.
+
+Higher positive correlation indicates stronger agreement between changes in the target property and changes in the generated property.
+
+### Maximum Mean Discrepancy — MMD
+
+MMD compares the distributions of generated and target property values using an RBF kernel.
+
+Lower MMD indicates closer distributional agreement.
+
+However, MMD should **not be interpreted in isolation**.
+
+A low MMD value does not necessarily imply superior generative performance. A compressed model may, for example, generate less diverse or structurally simpler outputs whose aggregate property distributions happen to resemble the target distribution.
+
+For this reason, MMD is interpreted jointly with:
+
+- chemical validity
+- MAE
+- RMSE
+- Pearson correlation
+- qualitative distribution analysis
+
+---
+
+## Distribution Analysis
+
+Selected molecular-property distributions can be visualized by comparing valid generated samples with the corresponding reference distribution.
+
+### Molecular Weight
+
+![Generated vs Test — Molecular Weight](MolWt_hist.png)
+
+### Heavy Atom Count
+
+![Generated vs Test — Heavy Atom Count](HeavyAtomCount_hist.png)
+
+### Ring Count
+
+![Generated vs Test — Ring Count](RingCount_hist.png)
+
+These plots complement the numerical metrics and help illustrate whether the generated structures reproduce important aspects of the target property distributions.
+
+---
+
+## Complete Experimental Pipeline
+
+```text
+Open Macromolecular Genome
+          Dataset
+             │
+             ▼
+     Data Preprocessing
+             │
+             ▼
+ ┌─────────────────────────┐
+ │ Polymer Representation  │
+ │           +             │
+ │ 15 Molecular Properties │
+ └────────────┬────────────┘
+              │
+              ▼
+   Base Transformer Model
+hkqiu/PolyTAO-BigSMILES_Version
+              │
+              ▼
+ Property-Conditioned Teacher
+              │
+              ▼
+     Teacher Hidden States
+              │
+              │
+              │ Knowledge Distillation
+              ▼
+ Property-Conditioned Student
+              │
+              ▼
+      Polymer Generation
+              │
+              ▼
+        RDKit Validation
+              │
+              ▼
+ Valid Generated Structures
+              │
+              ▼
+ RDKit Property Recalculation
+              │
+              ▼
+ ┌───────────────────────────┐
+ │ MAE                       │
+ │ RMSE                      │
+ │ MMD                       │
+ │ Pearson Correlation       │
+ │ Chemical Validity         │
+ └─────────────┬─────────────┘
+               │
+               ▼
+   Capacity–Performance
+          Analysis
+```
+
+---
+
+## Model Files
+
+The complete property-conditioned student model requires more than the standard Transformer checkpoint.
+
+The model repository may contain files such as:
+
+```text
+config.json
+generation_config.json
+model.safetensors
+tokenizer_config.json
+tokenizer files
+property_proj_student.pt
+```
+
+The standard Transformer weights are stored in:
+
+```text
+model.safetensors
+```
+
+The learned property-conditioning projection is stored separately in:
+
+```text
+property_proj_student.pt
+```
+
+This projection maps the 15-dimensional molecular-property vector into the Transformer hidden space.
+
+Because the property-conditioning mechanism is implemented outside the standard T5 forward interface, loading only the Transformer checkpoint is **not sufficient to reproduce property-conditioned generation**.
+
+---
+
+## Usage
+
+This model requires a custom property-conditioning step.
+
+It should therefore not be treated as a standard text-to-text T5 model that can reproduce the research pipeline using only a generic Hugging Face `pipeline()` call.
+
+The basic loading procedure is:
+
+```python
+import torch
+import torch.nn as nn
+
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+MODEL_ID = "FahimehBahman/PolyTAO-Student"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    MODEL_ID
+).to(device)
+
+model.eval()
+
+property_proj = nn.Linear(
+    15,
+    model.config.d_model
+).to(device)
+```
+
+The learned weights from:
+
+```text
+property_proj_student.pt
+```
+
+must then be loaded into `property_proj`.
+
+Property-conditioned generation requires the following steps:
+
+```text
+Polymer Input
+      │
+      ▼
+ Tokenization
+      │
+      ▼
+Transformer Encoder
+      │
+      ▼
+Encoder Hidden States
+      ▲
+      │
+15 Molecular Properties
+      │
+      ▼
+Property Projection
+      │
+      ▼
+Normalized Property Embedding
+      │
+      ▼
+Conditioned Encoder Representation
+      │
+      ▼
+Autoregressive Generation
+      │
+      ▼
+Generated Polymer
+```
+
+The complete training and evaluation implementation is provided in the accompanying PolyTAO source-code repository.
+
+---
+
+## Intended Use
+
+PolyTAO Student is intended primarily for research involving:
+
+- knowledge distillation
+- Transformer model compression
+- property-conditioned generation
+- polymer generative modeling
+- capacity–performance trade-offs
+- efficient generative Transformer architectures
+- evaluation of compressed generative models
+
+The checkpoint can also serve as a reproducible artifact for studying aggressive Transformer depth reduction under feature-level knowledge distillation.
+
+---
+
+## Limitations
+
+This model has several important limitations.
+
+- It is a **research prototype**, not a production model.
+- Chemical validity does not imply synthesizability.
+- Chemical validity does not imply stability or experimental feasibility.
+- Generated structures have not been experimentally validated.
+- RDKit-based validation is a computational validity check rather than experimental verification.
+- Property conditioning is limited to the molecular descriptors used during training.
+- Performance depends on preprocessing and decoding configuration.
+- Distributional similarity does not guarantee structural diversity.
+- MMD should not be interpreted independently of validity and other evaluation metrics.
+- Results from this model should not be used directly for safety-critical, medical, industrial, or experimental chemical decisions without independent expert validation.
+
+---
+
+## Research Context
+
+PolyTAO was developed as part of the M.Sc. thesis:
+
+**Knowledge-Distilled Transformers for Property-Conditioned Polymer Generation**
+
+The research investigates the effect of Transformer capacity reduction on property-conditioned polymer generation.
+
+The work combines:
+
+- Transformer-based generative modeling
+- molecular-property conditioning
+- feature-level knowledge distillation
+- model compression
+- computational chemistry descriptors
+- systematic quantitative evaluation
+
+Multiple student configurations were trained and evaluated to investigate how model depth affects generative behavior and property fidelity.
+
+The published checkpoint represents the **20% capacity student configuration**.
+
+---
+
+## Tools and Libraries
+
+The project uses:
+
+- Python
+- PyTorch
+- PyTorch Lightning
+- Hugging Face Transformers
+- RDKit
+- NumPy
+- Pandas
+- Scikit-learn
+- Matplotlib
+
+---
+
+## Author
+
+**Fahimeh Bahman**
+
+M.Sc. Data Science  
+Software Engineer & Machine Learning Researcher
+
+Research interests include:
+
+- Generative AI
+- Transformer Models
+- Knowledge Distillation
+- Model Compression
+- Machine Learning Systems
+- AI for Software Engineering