t3d-trajectory-distillation

name	t3d-trajectory-distillation
title	T3D: Few-Step Diffusion LMs via Trajectory Self-Distillation
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2602.12262
keywords	["Diffusion Language Models","Distillation","Few-Step Generation","Self-Supervision","Mode-Seeking"]
description	Accelerate diffusion language models from many to few generation steps using trajectory self-distillation. Collect clean-noisy state pairs along teacher trajectories, apply reverse-KL mode-seeking distillation, and weight losses by decoding order to prioritize early predictions where cascading errors compound.

T3D: Few-Step Diffusion LMs via Trajectory Self-Distillation

Problem Context

Diffusion language models require many sampling steps (25-50+) for high-quality generation. Few-step variants are needed for deployment, but naive training on random corruption schedules causes train-test mismatch: inference uses confidence-based schedules while training sees uniform random masking. T3D solves this by collecting trajectory pairs from teacher's actual inference schedule and using mode-seeking distillation.

Core Concept

T3D operates in two phases: (1) collect (clean, intermediate) state pairs along teacher trajectories during inference, (2) train student via reverse-KL objective that promotes mode-seeking (concentrating on high-probability outputs) with path-consistency weighting. This addresses the multimodal posterior problem where forward-KL spreads mass across modes.

Architecture Overview

• Trajectory collection: Run full teacher steps, save clean (x₀) and intermediate (xₜ) states
• Reverse-KL distillation: Use discriminative KL divergence promoting mode concentration
• Path-consistency weighting: Weight token losses by decoding order (early tokens > later tokens)
• Mode-seeking loss: Prevent mode-covering behavior in highly multimodal posteriors
• Few-step inference: Student generates quality outputs in 2-4 steps

Implementation

Step 1: Collect trajectory pairs from teacher

import torch
import torch.nn.functional as F
from typing import List, Tuple, Dict

class TrajectoryCollector:
    """Collect (clean, corrupted) state pairs from teacher inference."""

    def __init__(
        self,
        teacher_model,
        num_collection_steps: int = 50,
        corruption_schedule: str = 'linear'
    ):
        self.teacher_model = teacher_model
        self.num_collection_steps = num_collection_steps
        .corruption_schedule = corruption_schedule

     () -> torch.Tensor:
        
         .corruption_schedule == :
            
             torch.linspace(, , num_steps)
         .corruption_schedule == :
            
            t = torch.linspace(, , num_steps)
             torch.cos(t * torch.pi / )
        :
             ValueError()

     () -> torch.Tensor:
        
        num_mask = (sigma * (clean_tokens))
        mask_positions = torch.randperm((clean_tokens))[:num_mask]

        corrupted = clean_tokens.clone()
        corrupted[mask_positions] = mask_token_id

         corrupted

     () -> [[, torch.Tensor]]:
        
        trajectories = []
        corruption_schedule = .get_corruption_schedule(.num_collection_steps)

         traj_idx  (num_trajectories):
            
            clean_tokens = .teacher_model.generate(prompt, num_steps=)

            
            intermediates = []
             step_idx, sigma  (corruption_schedule):
                corrupted = .corrupt_tokens(
                    clean_tokens,
                    sigma.item(),
                    mask_token_id=.teacher_model.mask_token_id
                )
                intermediates.append(corrupted)

            trajectories.append({
                : clean_tokens,
                : intermediates,
                : corruption_schedule
            })

         trajectories