Haotian Sun

Matryoshka Agent: Unfolding Sub-Agents for Long-Horizon Machine Learning Engineering

2026-07-27T00:00:00+00:00

Machine learning engineering tasks demand extended decision-making cycles that stretch far beyond the horizon a single agent can reliably manage. Matryoshka Agent decomposes agentic problem solving into a coordinated hierarchy: an Orchestrator maintains high-level strategy and delegates, while Sub-Agents unfold to execute concrete, bounded tasks and report back. This separation keeps the Orchestrator’s context focused on strategy rather than execution detail, allowing the system to sustain long-horizon workflows. The framework enables Qwen3-4B-Instruct to reach Orchestrator performance comparable to o4-mini, and delivers up to 36.7% performance gains when paired with larger backbone models.

STACX: A Modular Infrastructure for End-to-End Agentic RL

2026-07-22T00:00:00+00:00

STACX (Stand Alone Complex) is a modular framework for agentic training, evaluation, and reinforcement learning. It standardizes the full pipeline — dataset config, rollout scheduler, agent and sandbox, verifier reward, trainer — so that new methods and new benchmarks slot into the same infrastructure. STACX supports training with GRPO, SFT, DAgger, and on-policy distillation; evaluating agents across benchmark environments; and integrating new environments. Three engines are coordinated under the hood: a task scheduler driving agents and scaffolds, a Docker-based sandbox environment engine, and Megatron/FSDP training backends. Agents can be implemented in-house, keeping the training loop fully visible, or installed as external black-box agents such as OpenHands and Terminus-2.

Forward-Free Diffusion Language Models with BPTT-Free Looped Refinement

2026-06-06T00:00:00+00:00

Diffusion language models conventionally depend on a hand-designed forward corruption process, which constrains what the model can learn to undo. We introduce FReDA, a diffusion-based language model that eliminates the prescribed forward process entirely. Instead, language modeling is framed as recursive distribution refinement, in which model-generated drafts serve as implicit intermediate states — and training avoids backpropagation through time across the refinement loop. Text is produced by iterative refinement, either through self-refinement or best-of-N selection among parallel candidates. At sub-8B scale, FReDA-4B achieves up to 15% gains on reasoning and coding tasks over larger diffusion models while delivering 1.5-1.8x inference speedups over diffusion baselines.

Exploration-Driven Optimization for Test-Time Large Language Model Reasoning

2026-05-11T00:00:00+00:00

Post-training combined with inference-time scaling significantly enhances the reasoning and alignment capabilities of large language models (LLMs). A fundamental tension arises, however: inference-time methods benefit from diverse sampling from a relatively flat probability distribution, whereas reinforcement learning (RL)-based post-training inherently sharpens these distributions. We propose Exploration-Driven Optimization (EDO), which extends reward-biasing-style exploration objectives to iterative post-training and integrates them into standard RL objectives, encouraging greater diversity in sampled solutions while enabling more effective inference-time computation. We incorporate EDO into iterative Direct Preference Optimization (iDPO) and Group Relative Policy Optimization (GRPO), yielding ED-iDPO and ED-GRPO, and show consistent gains in test-time reasoning across extensive experiments.

Spectral Representation-based Reinforcement Learning

2025-12-17T00:00:00+00:00

In real-world applications with large state and action spaces, reinforcement learning (RL) typically relies on function approximation to represent policies, value functions, and dynamics models. Although expressive approximators such as neural networks are powerful, they often present theoretical ambiguities, suffer from optimization instability and exploration difficulty, and incur substantial computational cost. We introduce the perspective of spectral representations as a solution to these difficulties. Stemming from the spectral decomposition of the transition operator, this framework yields an effective abstraction of the system dynamics for subsequent policy optimization while providing a clear theoretical characterization. We show how to construct spectral representations for transition operators with latent-variable or energy-based structure, each implying a different learning method and realizing an effective RL algorithm under the framework, and we provably extend the spectral view to partially observable MDPs. On more than 20 challenging tasks from the DeepMind Control Suite, the resulting algorithms match or surpass state-of-the-art model-free and model-based baselines.

Towards Better Instruction Following Retrieval Models

2025-05-27T00:00:00+00:00

Modern information retrieval (IR) models, trained exclusively on standard query–passage pairs, struggle to effectively interpret and follow explicit user instructions. We introduce InF-IR, a large-scale, high-quality training corpus tailored for enhancing retrieval models in instruction-following IR. InF-IR expands traditional training pairs into over 38,000 expressive instruction–query–passage triplets as positive samples. For each positive triplet, we generate two additional hard negatives by poisoning both instructions and queries, rigorously validated by an advanced reasoning model to ensure semantic plausibility while maintaining instructional incorrectness. Unlike existing corpora that primarily support computationally intensive reranking for decoder-only language models, the highly contrastive triplets in InF-IR also enable efficient representation learning for smaller encoder-only models, facilitating direct embedding-based retrieval. Using this corpus, we train InF-Embed, an instruction-aware embedding model optimized through contrastive learning and instruction–query attention. Across five instruction-based retrieval benchmarks, InF-Embed surpasses competitive baselines by 8.1% in p-MRR.

Matryoshka Pilot: Learning to Drive Black-Box LLMs with LLMs

2025-05-03T00:00:00+00:00

Matryoshka Pilot proposes a framework that uses lightweight LLMs to “drive” black-box LLMs on complex reasoning and agentic tasks. By decomposing control decisions across nested layers of increasing capability, the pilot models progressively refine prompts, tool calls, and plans handed to the frontier black-box model. The method achieves strong improvements in task success and cost efficiency across a suite of reasoning and decision-making benchmarks.

Spectral Representation for Causal Estimation with Hidden Confounders

2025-05-02T00:00:00+00:00

We propose a method for causal effect estimation in settings with hidden confounders, via spectral representations derived from proxy variables. By formulating the problem in a reproducing kernel Hilbert space and leveraging operator-theoretic tools, we derive tractable estimators that recover causal effects without requiring direct access to the confounder. Our approach generalizes instrumental variable and proxy-based methods, and we demonstrate both theoretical guarantees and strong empirical performance across synthetic and real-world benchmarks.

AmorLIP: Efficient Language-Image Pretraining via Amortization

2025-05-01T00:00:00+00:00

Contrastive Language-Image Pre-training (CLIP) has demonstrated strong zero-shot performance across diverse downstream text-image tasks. Existing CLIP methods typically optimize a contrastive objective using negative samples drawn from each minibatch. To achieve robust representation learning, these methods require extremely large batch sizes and escalate computational demands to hundreds or even thousands of GPUs. Prior approaches to mitigate this issue often compromise downstream performance, prolong training duration, or face scalability challenges with very large datasets. To overcome these limitations, we propose AmorLIP, an efficient CLIP pretraining framework that amortizes expensive computations involved in contrastive learning through lightweight neural networks, which substantially improves training efficiency and performance. Leveraging insights from a spectral factorization of energy-based models, we introduce novel amortization objectives along with practical techniques to improve training stability. Extensive experiments across 38 downstream tasks demonstrate the superior zero-shot classification and retrieval capabilities of AmorLIP, consistently outperforming standard CLIP baselines with substantial relative improvements of up to 12.24%.

MedAdapter: Efficient Test-Time Adaptation of Large Language Models Towards Medical Reasoning

2024-11-01T00:00:00+00:00

MedAdapter presents an efficient test-time adaptation method for applying large language models to medical reasoning tasks. Rather than finetuning the base LLM, MedAdapter trains a lightweight adapter module that operates at inference time, improving downstream accuracy on medical QA, diagnosis, and reasoning benchmarks while remaining compatible with black-box commercial LLMs.