cggr 0.4.3

Confidence-Gated Gradient Routing for Efficient Transformer Training

CGGR - Confidence-Gated Gradient Routing

[!WARNING] This is highly experimental, and prepare for the worst hope for the best.

Selective loss computation for Transformer training. Only hard tokens contribute to loss, providing actual backward pass savings.

Installation

pip install cggr

Requires CUDA + Triton.

Why CGGR?

Metric	Standard Training	CGGR (Batch Split)	Benefit
Backward Pass	100% of tokens	25% of tokens	4x cheaper backward pass
Forward Pass	1.0x cost	~1.1x cost (Pass 1 + 2)	Negligible overhead (~9ms)
Total Speed	1.0x (Baseline)	1.4x - 2.0x faster	Significant training acceleration
Data Efficiency	Learns from all tokens	Prioritizes hard tokens	Learns faster from hard examples
Memory	High (full graph)	Lower (sparse graph)	Can increase batch size

Benchmarks

Hardware: RTX 3060 Model: SmolLM-135M (Llama architecture)
Dataset: FineWeb-Edu

%%{init: {'theme': 'dark', 'themeVariables': { 'fontSize': '16px'}}}%%
xychart-beta
    title "Training Step Time (Lower is Better)"
    x-axis ["Standard", "CGGR (Optimized)"]
    y-axis "Time (ms)" 0 --> 350
    bar [309, 220]

Configuration	Forward (ms)	Backward (ms)	Total Step (ms)	Speedup
Standard Training	118 ms	185 ms	309 ms	1.0x
CGGR (Optimized)	127 ms	93 ms	220 ms	1.40x

Quick Start

1. Batch Splitting (Recommended)

The most efficient way to use CGGR is via CGGRModel. It uses a lightweight router to score difficulty and only computes gradients for hard tokens.

from cggr import CGGRModel, create_truncated_router
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("...").cuda()

# Create lightweight router (shares weights, 0 extra memory)
router = create_truncated_router(model, num_layers=4)

# Wrap model
cggr_model = CGGRModel(
    model, 
    router=router, 
    min_tokens_ratio=0.25
)

# Train
loss = cggr_model(input_ids, labels=labels)
loss.backward()

2. Manual Integration (CGGRLoss)

If you cannot use CGGRModel (e.g. specialized architectures), you can use CGGRLoss manually.

from cggr import CGGRLoss

criterion = CGGRLoss(
    scoring='combined',      # 'entropy', 'margin', 'loss', 'combined'
    selection='stratified',  # 'topk', 'stratified', 'sequence_aware'
    min_tokens_ratio=0.25,
    warmup_steps=1000,
)

for batch in dataloader:
    logits = model(input_ids)
    loss = criterion(logits, targets)  # Only hard tokens
    loss.backward()
    optimizer.step()
    criterion.step()

3. Native Integration (e.g. MoE/SRDE)

For architectures like SRDE that require static tensor shapes, you can use CGGRScorer to generate a routing mask instead of splitting the batch.

from cggr import CGGRScorer

# 1. Initialize Scorer
self.scorer = CGGRScorer(router, min_tokens_ratio=0.5)

# 2. Get Mask
difficulty, mask, info = self.scorer(input_ids)

# 3. Apply Mask (Null Routing)
# mask is Boolean: True=Hard (Route to Expert), False=Easy (Skip/Null)
expert_output = expert_layer(x) * mask.unsqueeze(-1)

Scoring Strategies

Strategy	Description	Best For
entropy	High entropy = hard	General training
margin	Small top-2 margin = hard	Classification
loss	High loss = hard	Direct optimization
combined	All signals combined	Best overall

Selection Strategies

Strategy	Description	Benefit
topk	Top-k hardest tokens	Simple, fast
stratified	Sample from difficulty buckets	Prevents forgetting
sequence_aware	Ensure coverage per sequence	Preserves structure

Dynamic Thresholding

Automatically adjusts token ratio based on batch confidence:

• Low confidence → more tokens (model is learning)
• High confidence → fewer tokens (model has converged)

CGGRLoss(dynamic_threshold=True, threshold_sensitivity=0.5)

Full API

CGGRLoss(
    # Scoring
    scoring='combined',
    
    # Selection
    selection='topk',
    num_strata=4,                  # For stratified
    min_tokens_per_sequence=1,     # For sequence_aware
    
    # Thresholding
    dynamic_threshold=True,
    threshold_sensitivity=0.5,
    
    # Curriculum
    min_tokens_ratio=0.25,
    warmup_steps=1000,
)

Performance

Config	Backward FLOPs	Overhead
Standard Loss	100%	0%
CGGR (25% tokens)	~25%	~0%