Task-Aligned Context Distillation: Set-Wise Selection with Structural and Geometric Priors

Long-context distillation aims to compress noisy inputs into minimal sufficient subsets while preserving downstream reasoning performance. Existing methods fall into two paradigms: point-wise methods score context units without modeling inter-chunk dependencies, while heuristic set-wise methods optimize diversity without grounding in solver feedback. Both are susceptible to semantic collapse, structure blindness, or solver mismatch under strict compression budgets. We reframe context distillation as a **solver-aligned set-wise** optimization problem, and introduce **TaCo**, a solver-in-the-loop policy learning framework. Solving this formulation via naive reinforcement learning (RL) exposes a fundamental challenge: continuous embeddings provide no inductive bias, yielding discontinuous reward surface and high reward variance that renders credit assignment intractable. By identifying two orthogonal sources of reward instability, TaCo overcomes this with two complementary priors: a **structural prior** preserves the internal integrity of each context unit, and a **geometric prior** promotes the diversity across selected units. Together, these priors substantially reduce the effective RL search space, enabling stable credit assignment under sparse solver feedback. On academic long-context benchmarks, TaCo outperforms SOTA compressors on unstructured text and even exceeds the full-context baseline on structured text. On an industrial benchmark with extreme length, high noise and mixed structure, TaCo achieves 98% and 102% of full-context F1 at 2.18% and 9.71% token retention respectively. Beyond aggregate performance, mechanistic diagnostics confirm that TaCo's three core components—the structural prior, geometric prior, and solver-aligned policy—orthogonally mitigate all three failure modes, validating the representational design of the framework.