Longllada Diffusion Context

LongLLaDA: Unlocking Long Context Capabilities in Diffusion LLMs

Core Concept

LongLLaDA investigates long-context performance in diffusion-based LLMs, discovering unique characteristics unavailable in autoregressive models. Diffusion LLMs maintain remarkably stable perplexity during context extrapolation due to bidirectional attention exposure to symmetric relative position ranges. Combined with NTK-based Rotary Position Embedding (RoPE) scaling, the method achieves 6x context expansion without training. A systematic analysis reveals both strengths (synthetic QA, stable extrapolation) and weaknesses (aggregation tasks).

Architecture Overview

• Bidirectional Attention: Diffusion LLMs exposed to [-T_train, T_train-1] position range vs. autoregressive [0, T_train-1]
• Stable Perplexity: Unlike autoregressive models that diverge quickly, diffusion maintains stable perplexity under extrapolation
• Local Perception: Sliding-window-like behavior during context extrapolation
• NTK-Based RoPE Scaling: Applies scaling transformation to base frequencies during inference
• Position Interpolation: Scale relative positions to fit within training range

Implementation

Step 1: Understand RoPE Scaling Theory

Implement Rotary Position Embedding (RoPE) with dynamic frequency scaling:

import torch
import numpy as np

class RoPEScaler:
    """
    Rotary Position Embedding with scaling for context extension.
    """
    def __init__(self, dim, base=10000):
        self.dim = dim
        self.base = base

        # Precompute inverse frequencies
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def apply_rope(self, x, position_ids, scaling_factor=1.0):
        """
        Apply RoPE with optional scaling.

        Args:
            x: [batch, seq_len, dim] embeddings
            position_ids: [batch, seq_len] position indices
            scaling_factor: scale position frequencies (1.0 = no scaling)
        """
        seq_len = x.shape[1]
        device = x.device

        # Scale inverse frequencies for NTK method
        inv_freq_scaled = self.inv_freq / scaling_factor

        # Compute angles: theta_m = m * theta_d, where theta_d = base^(-2d/dim)
        t = position_ids.float().unsqueeze(-1)  # [batch, seq_len, 1]
        freqs = torch.outer(t.squeeze(0), inv_freq_scaled)
        # [seq_len, dim/2]

        # Convert to complex representation
        emb = torch.cat([freqs, freqs], dim=-1)  # [seq_len, dim]

        # Apply rotation
        cos = emb.cos()
        sin = emb.sin()

        # Rotate embeddings
        x_rot = (x * cos) + (self._rotate_half(x) * sin)

        return x_rot

    @staticmethod
    def _rotate_half(x):
        """Rotate x by 90 degrees in complex plane"""
        x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
        return torch.cat([-x2, x1], dim=-1)

Step 2: Implement NTK-Based Scaling

Apply the NTK (Neural Tangent Kernel) scaling approach for position extrapolation:

def compute_ntk_scaling_factor(original_seq_len, target_seq_len):
    """
    Compute NTK-based scaling factor for RoPE.

    The idea: if we trained with max_position original_seq_len,
    scale frequencies to extend to target_seq_len.

    Args:
        original_seq_len: training context window size
        target_seq_len: desired inference context window

    Returns:
        scaling_factor: multiply inverse frequencies by this
    """
    # Linear interpolation ratio
    scaling_factor = target_seq_len / original_seq_len

    # NTK refinement: smooth scaling
    # scaling = alpha^(log(target/original) / log(base))
    # For simplicity: linear scaling
    return scaling_factor

Step 3: Comparative Analysis: Diffusion vs. Autoregressive

Understand why diffusion LLMs are more stable:

def analyze_position_exposure(model_type='diffusion', seq_len=512):
    """
    Analyze position range exposure during training.

    Diffusion LLMs: bidirectional attention → [-seq_len, seq_len]
    Autoregressive: causal attention → [0, seq_len]
    """
    if model_type == 'diffusion':
        # Bidirectional: every token can attend to all positions
        min_rel_pos = -(seq_len - 1)
        max_rel_pos = (seq_len - 1)
        pos_range = np.arange(min_rel_pos, max_rel_pos + 1)

        print(f"Diffusion LLM position range: [{min_rel_pos}, {max_rel_pos}]")
        print(f"Symmetric around 0: {-min_rel_pos == max_rel_pos}")

        return pos_range

    else:
        # Autoregressive: only attend to past
        min_rel_pos = 0
        max_rel_pos = (seq_len - 1)
        pos_range = np.arange(min_rel_pos, max_rel_pos + 1)

        print(f"Autoregressive position range: [{min_rel_pos}, {max_rel_pos}]")
        print(f"Asymmetric, only positive positions")

        return pos_range

Step 4: Apply NTK Scaling During Inference

Scale positions and frequencies to enable context extension:

class LongLLaDiffer:
    """
    Applies LongLLaDA: NTK-based RoPE scaling for diffusion LLMs.
    """
    def __init__(self, model, original_max_len=4096):
        self.model = model
        self.original_max_len = original_max_len
        self.rope_scaler = RoPEScaler(model.config.hidden_size)

    def extend_context(self, input_ids, max_new_len=24576):
        """
        Generate with extended context using NTK scaling.

        Args:
            input_ids: [batch, seq_len] token IDs (seq_len <= original_max_len)
            max_new_len: target context window (can be > original_max_len)

        Returns:
            output: [batch, seq_len] generated tokens
        """
        scaling_factor = compute_ntk_scaling_factor(
            self.original_max_len, max_new_len
        )

        # Get embeddings
        embeddings = self.model.get_input_embeddings()(input_ids)

        # Generate position IDs
        batch_size, seq_len = input_ids.shape
        position_ids = torch.arange(
            seq_len, device=input_ids.device
        ).unsqueeze(0).expand(batch_size, -1)

        # Apply RoPE with scaling
        embeddings = self.rope_scaler.apply_rope(
            embeddings,
            position_ids,
            scaling_factor=scaling_factor
        )

        # Forward through model with scaled embeddings
        outputs = self.model(inputs_embeds=embeddings)

        return outputs

    def evaluate_perplexity(self, test_texts, max_len=24576):
        """
        Evaluate perplexity with extended context.
        Compare to baseline (original max_len).
        """
        import math

        perplexities = {'baseline': [], 'extended': []}

        for text in test_texts:
            # Baseline: use original max length
            baseline_loss = self._compute_loss(
                text, max_len=self.original_max_len
            )
            baseline_ppl = math.exp(baseline_loss)
            perplexities['baseline'].append(baseline_ppl)

            # Extended: use NTK scaling
            extended_loss = self._compute_loss(
                text, max_len=max_len, use_scaling=True
            )
            extended_ppl = math.exp(extended_loss)
            perplexities['extended'].append(extended_ppl)

        return perplexities

    def _compute_loss(self, text, max_len, use_scaling=False):
        """Compute loss on text"""
        tokens = self.model.tokenizer.encode(text)[:max_len]
        input_ids = torch.tensor([tokens], device=self.model.device)

        with torch.no_grad():
            if use_scaling:
                outputs = self.extend_context(input_ids, max_new_len=max_len)
            else:
                outputs = self.model(input_ids)

            logits = outputs.logits
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = input_ids[..., 1:].contiguous()

            loss = torch.nn.functional.cross_entropy(
                shift_logits.view(-1, shift_logits.shape[-1]),
                shift_labels.view(-1)
            )

        return loss.item()

Step 5: Benchmark on Long-Context Tasks

Evaluate on retrieval, aggregation, and synthetic tasks:

def benchmark_long_context(model, max_len=24576):
    """
    Benchmark diffusion LLM on various long-context tasks.
    """
    benchmarks = {}

    # Task 1: Needle-In-A-Haystack (NIAH)
    print("Evaluating NIAH...")
    niah_scores = evaluate_needle_in_haystack(model, max_len=max_len)
    benchmarks['niah'] = niah_scores

    # Task 2: Retrieval
    print("Evaluating retrieval tasks...")
    retrieval_scores = evaluate_retrieval(model, max_len=max_len)
    benchmarks['retrieval'] = retrieval_scores

    # Task 3: Aggregation (harder for local perception)
    print("Evaluating aggregation...")
    aggregation_scores = evaluate_aggregation(model, max_len=max_len)
    benchmarks['aggregation'] = aggregation_scores

    # Task 4: Synthetic QA
    print("Evaluating synthetic QA...")
    qa_scores = evaluate_synthetic_qa(model, max_len=max_len)
    benchmarks['qa'] = qa_scores

    # Print summary
    print("\nBenchmark Results:")
    for task, scores in benchmarks.items():
        avg_score = np.mean(list(scores.values()))
        print(f"  {task}: {avg_score:.2%}")

    return benchmarks

def evaluate_needle_in_haystack(model, max_len, num_tests=100):
    """
    Evaluate ability to find specific fact in long context.
    """
    scores = {'found': 0, 'not_found': 0}

    for _ in range(num_tests):
        # Create haystack: mostly irrelevant text
        haystack = "This is irrelevant context. " * (max_len // 10)

        # Insert needle: important fact
        needle = "The capital of France is Paris."
        haystack = haystack[:max_len//2] + needle + haystack[max_len//2:]

        # Query
        query = "What is the capital of France?"

        # Generate answer
        answer = model.generate_with_context(haystack, query)

        # Check if found
        if "Paris" in answer:
            scores['found'] += 1
        else:
            scores['not_found'] += 1

    return scores

Practical Guidance

• Scaling Factor: Start with target_len / original_len; can experiment with polynomial scaling
• Position Interpolation: For diffusion LLMs, linear scaling works well due to bidirectional training
• Testing Tasks: Always evaluate on NIAH, retrieval, and aggregation to understand model behavior
• Stability Validation: Monitor perplexity across different context lengths; should remain stable
• Integration: Modify RoPE scaling at inference only; no model retraining required
• Comparison: Benchmark against Position Interpolation and other extrapolation methods

Reference

Paper: arXiv:2506.14429 Key metrics: 6x context expansion (24k tokens), stable perplexity, strong NIAH performance Architecture differences: Bidirectional attention enables symmetric position exposure Related work: Position interpolation, RoPE scaling, context extension methods