ARTN / README.md
kishxrx's picture
Upload README.md with huggingface_hub
bece1fb verified
|
Raw History Blame Contribute Delete
11.2 kB
metadata
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

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)

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:

@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


Last Updated: October 15, 2026

Version: v0.2-research-progress