AriaLM / docs /HLD.md
krishnah27's picture
Upload folder using huggingface_hub
30e9297 verified
|
Raw
History Blame Contribute Delete
3.94 kB
# High-Level Design (HLD)
## System Overview
Music Generation LLM is a symbolic music generation system that learns patterns from
classical MIDI music and generates new compositions. It operates entirely on symbolic
note representations (MIDI events), not raw audio waveforms.
## Architecture Diagram
```
β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
β”‚ USER INTERFACE β”‚
β”‚ CLI: train / generate / both β”‚
β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
β”‚
β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
β”‚ ORCHESTRATOR (s00_main) β”‚
β”‚ Parses args, wires components β”‚
β””β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
β”‚ β”‚ β”‚
β”Œβ”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”€β”
β”‚ DATA β”‚ β”‚ MODEL β”‚ β”‚ GENERATION β”‚
β”‚ PIPELINE β”‚ β”‚ PIPELINE β”‚ β”‚ PIPELINE β”‚
β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€
β”‚ HF Download β”‚ β”‚ MusicTrans- β”‚ β”‚ Autoregressive β”‚
β”‚ MIDI Parse β”‚ β”‚ former Build β”‚ β”‚ w/ KV-cache β”‚
β”‚ REMI Token β”‚ β”‚ Train Loop β”‚ β”‚ Top-p/k Sample β”‚
β”‚ DataLoader β”‚ β”‚ AMP + Ckpt β”‚ β”‚ MIDI Export β”‚
β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
```
## Component Responsibilities
### Data Pipeline (s02 + s03)
- Download 4,796 MIDI files from HuggingFace
- Parse with `pretty_midi` library
- Tokenize using REMI scheme (Note On/Off, Velocity, TimeShift, Bar, Position)
- Create PyTorch DataLoaders with padding and random cropping
### Model (s04)
- LLaMA-style Transformer (6 layers, 256 dim, 8 heads)
- Grouped Query Attention with 4 KV heads (50% memory reduction)
- SwiGLU feed-forward with RMSNorm
- RoPE for positional encoding
- Gradient checkpointing support
### Training (s05)
- AdamW optimizer with cosine LR warmup
- Mixed precision (BF16/FP16) via PyTorch AMP
- Gradient accumulation (effective batch = 32)
- Early stopping with patience
- TensorBoard logging
### Generation (s06)
- Autoregressive decoding with KV-cache
- Temperature scaling + Top-k + Top-p (nucleus) sampling
- Repetition penalty for diverse output
- Direct MIDI file export
## Data Flow
```
MIDI Files β†’ REMI Tokens β†’ Training β†’ Trained Model β†’ Generation β†’ MIDI Output
```
## Constraints & Decisions
| Decision | Rationale |
|------------------------------|----------------------------------------------|
| Symbolic (MIDI) not audio | 100x smaller data, trainable on consumer GPU |
| GQA over standard MHA | 50% KV-cache memory reduction |
| Gradient checkpointing | ~50% memory savings, ~20% speed cost |
| BF16 mixed precision | 50% memory, maintains numeric stability |
| REMI tokenization | Best published results for symbolic music |
| 5M params (not 100M+) | Fits in <=4GB VRAM, fast training |