Hugging Face's logo

kishxrx
/
ARTN

ternary-quantization
1.58-bit
experimental
language-model
tiny-stories
adafactor
bitlinear
adamw
Model card Files
xet
 Community
ARTN / README.md
kishxrx's picture
kishxrx
Upload README.md with huggingface_hub
bece1fb verified
preview code
|
Raw History Blame Contribute Delete
11.2 kB
---
license: mit
tags:
- ternary-quantization
- 1.58-bit
- experimental
- language-model
- tiny-stories
- adafactor
- bitlinear
- adamw
license_name: mit
license_link: https://opensource.org/licenses/MIT
---
# ARTN - Adaptive Ternary Recurrent Network
## 1.58-bit Ternary Quantization for Language Models
**Experimental implementation of ternary (1.58-bit) quantization for training small language models on extremely limited hardware.**
**Status: EXPERIMENTAL - Research in Progress**
---
## Project Overview
This project explores whether **ternary quantization** (weights constrained to -1, 0, +1) can be used to train small language models on hardware with:
| Constraint | Value |
|------------|-------|
| RAM | 4GB |
| CPU | 2-core (no GPU) |
| Storage | ~20GB |
---
## Model Architecture
| Parameter | Value |
|-----------|-------|
| Model type | BitLinearGPT |
| Parameters | 6,856,608 (~6.8M) |
| Hidden size (d_model) | 256 |
| Layers | 6 |
| Attention heads | 8 |
| Feed-forward size | 1024 |
| Max sequence length | 256 |
| Vocabulary size | 4,000 |
| Quantization | Ternary (1.58-bit) |
---
## Research Timeline
### Phase 1: Initial Experiments (August 2026)
Early experiments with Adafactor optimizer:
| Attempt | Config | Result | Issue |
|---------|--------|--------|-------|
| 5k steps | Beta1=None | COLLAPSED | Model fell into repetitive patterns |
| 50k steps | Beta1=None | COLLAPSED | Entropy dropped to 8.2% |
| 50k steps | Beta1=0.9 | HEALTHY | Learning too slow, val loss 7.56 |
**Key Finding:** Beta1=None caused collapse. Beta1=0.9 too conservative.
---
### Phase 2: Root Cause Investigation (September 2026)
**September 26, 2026:** Critical discovery - Adafactor implementation bug identified.
| Issue | Details |
|-------|---------|
| Bug location | Original Adafactor implementation |
| Bug type | Variance computed but not applied to updates |
| Impact | Optimizer behaved like SGD with momentum |
| Evidence | Lines 152, 170 in `src/training/adafactor.py` |
**Verification:** Original optimizer state matched broken implementation, not canonical Adafactor.
---
### Phase 3: Working Models (September 2026)
**After identifying the bug, switched to AdamW optimizer:**
#### 1. FP32 + AdamW Baseline (August 19, 2026)
| Metric | Value |
|--------|-------|
| Optimizer | AdamW |
| Steps | 5,000 |
| Val Loss | **1.69** |
| Perplexity | **5.44** |
| Generation | ✅ Coherent stories |
| Status | **SUCCESS - Baseline established** |
**File:** `checkpoint-fp32-adamw-working.pt`
---
#### 2. BitLinear + AdamW (September 10-11, 2026)
| Metric | Value |
|--------|-------|
| Optimizer | AdamW |
| Steps | 9,500 |
| Val Loss | **1.92** |
| Perplexity | **6.81** |
| Generation | ✅ Coherent stories |
| Quantization | Ternary (1.58-bit) |
| Status | **SUCCESS - Ternary quantization works** |
**File:** `checkpoint-bitlinear-adamw-working.pt`
---
### Phase 4: Canonical Adafactor Investigation (October 15, 2026)
Implemented and tested canonical Adafactor (matching PyTorch reference):
| Metric | Value |
|--------|-------|
| Implementation | Verified against PyTorch (0.025% accuracy) |
| Val Loss | 4.82 |
| Perplexity | 124.10 |
| Optimizer State Memory | 0.16 MB (316x less than AdamW) |
| Observed Performance | 2.0x gap vs AdamW on tested configurations |
**Note:** Performance gap observed, but configurations differed (weight decay, momentum). Investigation ongoing.
---
## Key Findings
### 1. Ternary Quantization is VIABLE
| Evidence | Result |
|----------|--------|
| BitLinear + AdamW | Val loss 1.92 (vs FP32 1.69) |
| Performance gap | 13% (acceptable for 85% memory reduction) |
| Generation quality | ✅ Coherent stories |
**Conclusion:** Ternary quantization works with proper optimizer (AdamW).
---
### 2. Optimizer Choice is CRITICAL
| Optimizer | Result | Reason |
|-----------|--------|--------|
| Adafactor (buggy) | ❌ FAILED | Variance not applied |
| Adafactor (canonical) | ⚠️ 2x gap | Configuration differences |
| **AdamW** | ✅ **WORKS** | Proven convergence |
---
### 3. Beta1 Trade-off in Adafactor
| Beta1 Value | Learning Speed | Stability | Result |
|-------------|---------------|-----------|--------|
| None (0.0) | FAST | UNSTABLE | COLLAPSED |
| 0.9 | SLOW | STABLE | HEALTHY but slow learning |
---
### 4. Memory vs Performance Trade-off
| Optimizer | Optimizer State | Performance | Trade-off |
|-----------|----------------|-------------|-----------|
| AdamW | 51.68 MB | ✅ Best (val loss 1.69-1.92) | More memory |
| Canonical Adafactor | 0.16 MB | ⚠️ Gap (val loss 4.82) | 316x less memory |
**Use case dependent:** AdamW for quality, Adafactor for memory-constrained scenarios.
---
## Checkpoints Included
### Working Models (Upload as evidence)
| File | Type | Val Loss | Status | Description |
|------|------|----------|--------|-------------|
| `checkpoint-fp32-adamw-working.pt` | FP32 | **1.69** | ✅ WORKING | AdamW baseline, coherent generation |
| `checkpoint-bitlinear-adamw-working.pt` | BitLinear | **1.92** | ✅ WORKING | Ternary quantization, coherent generation |
---
### Experimental/Failed Models (Already uploaded)
| File | Type | Val Loss | Status | Description |
|------|------|----------|--------|-------------|
| `checkpoint-step5000-experiment.pt` | BitLinear | 8.45 | ✅ Pipeline verified | First test, proved architecture works |
| `checkpoint-step50000-failed-collapsed.pt` | BitLinear | 3.16 | ❌ COLLAPSED | Beta1=None, repetitive output |
| `checkpoint-step50000-experiment-gibberish.pt` | BitLinear | 7.56 | ⚠️ Not useful | Beta1=0.9, healthy but slow learning |
---
## Training Configuration
### Working Models Configuration
| Setting | FP32 + AdamW | BitLinear + AdamW |
|---------|--------------|-------------------|
| Dataset | TinyStories | TinyStories |
| Steps | 5,000 | 9,500 |
| Batch size | 4 | 4 |
| Gradient accumulation | 8 | 8 |
| Effective batch size | 32 | 32 |
| Optimizer | AdamW | AdamW |
| Max Learning Rate | 0.001 | 0.001 |
| Min Learning Rate | 0.0001 | 0.0001 |
| Warmup steps | 100 | 100 |
| Weight decay | 0.01 | 0.01 |
| Gradient clipping | 1.0 | 1.0 |
---
### Failed Experiments Configuration
| Setting | Value |
|---------|-------|
| Dataset | TinyStories |
| Samples | 50,000 |
| Batch size | 4 |
| Gradient accumulation | 8 |
| Effective batch size | 32 |
| Optimizer | Adafactor (buggy) |
| Max Learning Rate | 0.0003 |
| Min Learning Rate | 0.00003 |
| Warmup steps | 1,000 |
| Weight decay | 0.01 |
| Gradient clipping | 1.0 |
---
## Total Time Invested
| Phase | Duration | Outcome |
|-------|----------|---------|
| Initial experiments (August) | 167 hours | Identified beta1 issues |
| Root cause investigation (September) | 5 days | Found Adafactor bug |
| Working models (August-September) | 15+ hours | ✅ Both models working |
| Canonical Adafactor (October) | 1 day | Implementation verified, performance gap observed |
---
## Lessons Learned
### Technical Lessons
1. **Optimizer implementation bugs are subtle** - Adafactor computed variance but didn't apply it
2. **Always verify against reference implementations** - PyTorch Adafactor differs from custom implementation
3. **AdamW is robust for ternary quantization** - Proven to work where Adafactor failed
4. **Memory vs performance trade-offs exist** - 316x memory savings vs 2x performance gap
5. **Lower loss ≠ better model** - Collapsed model had lower loss (3.16 vs 7.56) due to repetitive output
### Research Lessons
1. **Entropy monitoring is CRITICAL** - Loss alone doesn't indicate model health
2. **Systematic investigation pays off** - Line-by-line code audit found the bug
3. **Document everything** - 59 phases of investigation documented
4. **Test one variable at a time** - Controlled experiments essential for debugging
---
## How to Use Working Models
### FP32 + AdamW Baseline
```python
import torch
from src.model.gpt import GPT
from src.tokenizer.bpe_tokenizer import BPETokenizer
# Load model architecture (FP32)
model = GPT(
vocab_size=4000,
d_model=256,
n_heads=8,
n_layers=6,
d_ff=1024,
max_seq_len=256
)
# Load checkpoint
checkpoint = torch.load('checkpoint-fp32-adamw-working.pt', map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
# Load tokenizer
tokenizer = BPETokenizer.load('tokenizer.json')
# Generate text
input_ids = torch.tensor([tokenizer.encode("Once upon a time")])
# ... generation code ...
```
---
### BitLinear + AdamW (Ternary)
```python
import torch
from src.model.bitlinear import BitLinearGPT
from src.tokenizer.bpe_tokenizer import BPETokenizer
# Load model architecture (BitLinear)
model = BitLinearGPT(
vocab_size=4000,
d_model=256,
n_heads=8,
n_layers=6,
d_ff=1024,
max_seq_len=256
)
# Load checkpoint
checkpoint = torch.load('checkpoint-bitlinear-adamw-working.pt', map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
# Load tokenizer
tokenizer = BPETokenizer.load('tokenizer.json')
# Generate text
input_ids = torch.tensor([tokenizer.encode("Once upon a time")])
# ... generation code ...
```
---
## Tokenizers
| File | Description |
|------|-------------|
| `tokenizer-5k-experiment.json` | Tokenizer for early experiments |
| `tokenizer-50k-experiments.json` | Tokenizer for 50k step experiments |
---
## Source Code
Full source code available at GitHub:
**https://github.com/kishxrx/ARTN**
Includes:
- Complete training scripts
- Model architecture (FP32 and BitLinear)
- Optimizer implementations
- Training history documentation
- Monitoring tools
- 59 phases of investigation documented
---
## Documentation
Complete investigation documentation:
**https://github.com/kishxrx/ARTN/blob/main/experiments/results/ROOT_CAUSE_INVESTIGATION_SUMMARY.md**
Covers:
- Phases 1-59 of systematic debugging
- Root cause analysis
- Verification experiments
- Evidence chain
---
## Next Steps
| Phase | Status | Plan |
|-------|--------|------|
| 1 | ✅ COMPLETE | FP32 baseline established (val loss 1.69) |
| 2 | ✅ COMPLETE | BitLinear + AdamW working (val loss 1.92) |
| 3 | ✅ COMPLETE | Root cause identified (Adafactor bug) |
| 4 | ✅ COMPLETE | Canonical Adafactor implemented and tested |
| 5 | 🔄 IN PROGRESS | Scale training with AdamW (longer runs) |
| 6 | 📋 PLANNED | Evaluate on larger datasets |
| 7 | 📋 PLANNED | Release v1.0 with working models |
---
## Citation
If you use this work for learning or research:
```bibtex
@misc{artn-ternary-2026,
author = {kishxrx},
title = {ARTN: Adaptive Ternary Recurrent Network - 1.58-bit Quantization Experiments},
year = {2026},
publisher = {HuggingFace},
url = {https://huggingface.co/kishxrx/ARTN}
}
```
---
## License
MIT License
---
## Disclaimer
**THIS IS EXPERIMENTAL SOFTWARE.**
- Working models produce coherent text but are experimental
- Failed models are for learning purposes only
- Not suitable for production use without further testing
- No guarantees of any kind
---
## Contact
- GitHub: https://github.com/kishxrx/ARTN
- HuggingFace: https://huggingface.co/kishxrx
---
**Last Updated:** October 15, 2026
**Version:** v0.2-research-progress