alea-institute/kl3m-006-170m-checkpoint-57000

KL3M 170M, 6th Gen Model, 57K Checkpoint (Phase 2)

A 170M parameter language model trained on multi-domain legal text using the Muon optimizer with spectral clamping. This checkpoint represents Phase 2 multi-domain training following Phase 1 single-domain training (37K steps).

Model Details

• Architecture: Llama-based with Grouped Query Attention (GQA)
• Parameters: 181.7M (170M non-embedding)
• Training Steps: 57,000 (37K Phase 1 + 20K Phase 2)
• Tokens Processed: 15.53 billion
• Sequence Length: 4,096 tokens
• Precision: BF16
• Optimizer: Muon with spectral regularization (max condition: 3000)

Model Architecture

• Hidden Size: 576
• Layers: 30
• Attention Heads: 9 (3 KV heads with GQA)
• Intermediate Size: 1536
• Vocabulary: 131,072 tokens
• RoPE Theta: 100,000

Training Configuration

Phase 2 Multi-Domain Dataset

• Source: alea-institute/kl3m-data-sample-004-balanced
• Type: Multi-domain legal corpus
  • RECAP (Court filings): 48-52%
  • GovInfo (Federal regulations): 27-32%
  • USPTO (Patents): 3-5%
  • EDGAR (SEC filings): 5-7%
• Format: Streaming, shuffled with buffer=32
• Domain Quality: Balanced for broad legal knowledge

Optimizer (Muon)

• Muon Learning Rate: 7.30e-5 (depth-scaled from 1e-4)
• Auxiliary Learning Rate: 9e-5
• Muon Weight Decay: 1e-4
• Auxiliary Weight Decay: 1e-3
• Muon Momentum: 0.95
• Muon NS Steps: 3 (faster convergence vs Phase 1's 5)
• Batch Size: 6 per device
• Gradient Accumulation: 2 steps (effective batch: 12)
• Warmup Steps: 0 (continuing from Phase 1)
• LR Scheduler: Cosine with min ratio 0.1

Regularization

• Spectral Clamping:
  • Enabled on q_proj, o_proj, and lm_head
  • Max condition number: 3000 (relaxed from Phase 1's 2000)
  • Sigma floor: 5e-5 (relaxed from Phase 1's 1e-4)
  • Applied every 160 steps (relaxed from Phase 1's 10 steps)
  • Rationale: Relaxed constraints allow faster exploration of multi-domain manifolds
• Adaptive Gradient Clipping: Enabled (β=0.9, coeff=2.0)
• Label Smoothing: 0.01
• Entropy Regularization:
  • Entropy bonus weight: 0.003
  • Entropy target: 6.5 bits (weight: 0.003)
  • Activation norm weight: 0.0006
  • Loss chunk size: 1024 tokens

Training Infrastructure

• Mixed Precision: BF16
• Gradient Checkpointing: Enabled (non-reentrant)
• Flash Attention: Auto-enabled
• TF32 Mode: Auto

Spectral Health (Step 57K)

Analysis of weight matrix conditioning shows acceptable manifold quality with expected degradation from relaxed regularization:

Condition Numbers

• Attention Layers:
  • Median: 2307 (↑ 42% from Phase 1's 1620)
  • Mean: ~2500
  • P95: ~3200
  • Max: 3632 (↑ 46% from Phase 1's 2481)
• MLP Layers:
  • Median: ~5-10 (healthy)
  • Mean: ~8-12
  • Max: ~15-20 (excellent conditioning)
• LM Head: ~300-400 (good)

Singular Values

• Smallest σ_min: ~3.42e-4 (above σ_floor of 5e-5)
• Worst-conditioned layers:
  • Attention projection layers (q_proj, o_proj) approaching ceiling
  • K/V projections moderately conditioned

Key Finding: Attention layers show higher condition numbers than Phase 1 due to intentionally relaxed spectral regularization. This trade-off enables faster multi-domain learning and exploration while maintaining numerical stability.

Training Dynamics (Phase 2: Steps 37K-57K)

• Loss: Competitive across all domains
• Gradient Norm: Stable with adaptive clipping
• Learning Rate: Gradual cosine decay
• Multi-Domain Quality: Balanced performance on RECAP, GovInfo, USPTO, EDGAR

Performance Comparison vs Phase 1

Metric	Phase 1 (37K)	Phase 2 (57K)	Change
Attention Conditioning	Median: 1620	Median: 2307	↑ 42%
Generation Quality	5.7/10	8.5/10	↑ 49%
Multi-Domain Coverage	EDGAR only	4 domains	✓
Spectral Clamp Freq	Every 10 steps	Every 160 steps	16x relaxed

Result: Intentional conditioning trade-off yields significantly better generation quality and multi-domain capability.

Generation Quality

Generates coherent, fluent legal text across multiple domains with no repetition issues. Significantly improved over Phase 1:

• Court documents (RECAP): Natural legal argumentation
• Federal regulations (GovInfo): Proper regulatory structure
• Patents (USPTO): Technical claim language
• SEC filings (EDGAR): Financial/corporate disclosure style

Suitable for multi-domain legal content generation, analysis, and understanding tasks.

Usage

from transformers import pipeline

# Create text generation pipeline
generator = pipeline(
    "text-generation",
    model="alea-institute/kl3m-006-170m-checkpoint-57000",
    torch_dtype="auto",
    device_map="auto"
)

# Generate text
outputs = generator(
    "This Agreement is entered into as of",
    max_new_tokens=100,
    temperature=0.7,
    top_p=0.95,
    repetition_penalty=1.15
)

print(outputs[0]['generated_text'])

Training Philosophy

This checkpoint demonstrates the intentional conditioning vs quality trade-off:

• Phase 1: Tight spectral regularization → Excellent conditioning, limited quality
• Phase 2: Relaxed regularization → Moderate conditioning, excellent quality

The 16x reduction in spectral clamping frequency allows the model to explore richer multi-domain manifolds while maintaining acceptable numerical stability (max condition < 4000).

Model Card Authors

Alea Institute

Citation

For technical details, see the paper: https://arxiv.org/abs/2504.07854

@misc{kl3m2025,
  title={KL3M: Knowledge-Guided Language Model Training},
  author={Alea Institute},
  year={2025},
  url={https://arxiv.org/abs/2504.07854},
  note={Trained with Muon optimizer and multi-domain spectral regularization}
}

License

Apache 2.0