Second-Order, First-Class: A Composable Stack for Curvature-Aware Training

• URL: http://arxiv.org/abs/2603.25976v1
• Date: Thu, 26 Mar 2026 23:36:37 GMT
• Title: Second-Order, First-Class: A Composable Stack for Curvature-Aware Training
• Abstract summary: We introduce Somax, a composable Optax-native stack that treats curvature-aware training as a single JIT-compiled step governed by a static plan.<n>Somax exposes first-class modules behind a single step interface and composes with Optax by applying standard gradient transformations.<n>We report system-oriented ablations showing that composition choices materially affect scaling behavior and time-to-accuracy.
• Score: 1.2891210250935148
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Second-order methods promise improved stability and faster convergence, yet they remain underused due to implementation overhead, tuning brittleness, and the lack of composable APIs. We introduce Somax, a composable Optax-native stack that treats curvature-aware training as a single JIT-compiled step governed by a static plan. Somax exposes first-class modules -- curvature operators, estimators, linear solvers, preconditioners, and damping policies -- behind a single step interface and composes with Optax by applying standard gradient transformations (e.g., momentum, weight decay, schedules) to the computed direction. This design makes typically hidden choices explicit and swappable. Somax separates planning from execution: it derives a static plan (including cadences) from module requirements, then runs the step through a specialized execution path that reuses intermediate results across modules. We report system-oriented ablations showing that (i) composition choices materially affect scaling behavior and time-to-accuracy, and (ii) planning reduces per-step overhead relative to unplanned composition with redundant recomputation.

Related papers

• Reconfiguration-Complete Motion Primitives with Constructive Planning for Deformable Planar Modular Robots [12.157600704536533]
  This letter introduces a square-cell abstraction that maps deformable rhombus modules to fixed-size grid cells.<n>We prove that every non-straight edge-connected configuration with $N geq 7$ can be transformed to a fixed staircase canonical using only admissible primitive motions.
  arXiv  Detail & Related papers  (2026-08-18T03:34:54Z)
• Scalable Deep Unfolding of Conic Optimizers [16.022960827366067]
  Deep unfolding (DU) accelerates iteratives by introducing learnable components and training them through unrolled iterations.<n>Unrolling a full-update conic solver such as COSMO exposes two obstacles that prior work on learned conic solvers has not.<n>We address the first obstacle with a matrix-free implicit differentiation rule, reducing memory from $O(n2)$ to $O(n)$ and enabling backpropagation at scales where direct factorization runs out of memory.
  arXiv  Detail & Related papers  (2026-06-11T18:58:09Z)
• Planning-aligned Token Compression for Long-Context Autonomous Driving [95.59023657139208]
  We propose a planning-aligned working memory framework built on conditional VQ-VA.<n>Compression is conditioned on both historical trajectory and a learned planning intent.<n>We evaluate on high-signal dynamic scenarios where historical context is most critical for behavior correctness.
  arXiv  Detail & Related papers  (2026-06-05T17:16:21Z)
• ReCAPA: Hierarchical Predictive Correction to Mitigate Cascading Failures [11.399289646886968]
  We propose a framework that uses prediction and contrast to adjust deviations across three levels: actions, subgoals, and trajectories.<n>We introduce two new metrics to quantify error propagation and recovery processes in tasks.<n>Experiments show that ReCAPA achieves competitive results on embodied agent benchmarks such as VisualAgentBench, MineDojo, and AI2-THOR.
  arXiv  Detail & Related papers  (2026-04-23T02:57:50Z)
• PRISM: Parallel Residual Iterative Sequence Model [52.26239951489612]
  We propose PRISM (Parallel Residual Iterative Sequence Model) to resolve this tension.<n>PRISM introduces a solver-inspired inductive bias that captures key structural properties of multi-step refinement in a parallelizable form.<n>We prove that this formulation achieves Rank-$L$ accumulation, structurally expanding the update manifold beyond the single-step Rank-$1$ bottleneck.
  arXiv  Detail & Related papers  (2026-02-11T12:39:41Z)
• Higher-order Linear Attention [59.92962330635185]
  quadratic cost of scaled dot-product attention is a central obstacle to scaling autoregressive language models to long contexts.<n>We introduce Higher-order Linear Attention (HLA), a causal, streaming mechanism that realizes higher interactions via compact prefix sufficient statistics.
  arXiv  Detail & Related papers  (2025-10-31T07:54:37Z)
• Kad: A Framework for Proxy-based Test-time Alignment with Knapsack Approximation Deferral [6.949966663998242]
  Large language models (LLM) still require further alignment to adhere to downstream task requirements and stylistic preferences.<n>As LLMs continue to scale in terms of size, the computational cost of alignment procedures increase prohibitively.<n>We propose a novel approach to circumvent these costs via proxy-based test-time alignment.
  arXiv  Detail & Related papers  (2025-10-30T21:38:45Z)
• Logits Replay + MoClip: Stabilized, Low-Cost Post-Training with Minimal Forgetting [6.653834890554154]
  We introduce Logits Replay + MoClip, a framework that compresses supervision in the logit space and stabilizes optimization at the update level.<n> Empirically, our method improves domain performance on Communication Technology tasks while mitigating forgetting on general benchmarks.
  arXiv  Detail & Related papers  (2025-10-10T08:55:32Z)
• PT$^2$-LLM: Post-Training Ternarization for Large Language Models [52.4629647715623]
  Large Language Models (LLMs) have shown impressive capabilities across diverse tasks, but their large memory and compute demands hinder deployment.<n>We propose PT$2$-LLM, a post-training ternarization framework tailored for LLMs.<n>At its core is an Asymmetric Ternary Quantizer equipped with a two-stage refinement pipeline.
  arXiv  Detail & Related papers  (2025-09-27T03:01:48Z)
• Constrained Entropic Unlearning: A Primal-Dual Framework for Large Language Models [14.321060805197874]
  Large Language Models (LLMs) deployed in real-world settings increasingly face the need to unlearn sensitive, outdated, or proprietary information.<n>Existing unlearning methods formulate forgetting and retention as a regularized trade-off, combining both objectives into a single scalarized loss.<n>We propose a new formulation of LLM unlearning as a constrained optimization problem: forgetting is enforced via a novel logit-margin flattening loss.
  arXiv  Detail & Related papers  (2025-06-05T17:55:23Z)
• Train with Perturbation, Infer after Merging: A Two-Stage Framework for Continual Learning [57.514786046966265]
  We propose textbfPerturb-and-Merge (P&M), a novel continual learning framework that integrates model merging into the CL paradigm to mitigate forgetting.<n>Our proposed approach achieves state-of-the-art performance on several continual learning benchmark datasets.
  arXiv  Detail & Related papers  (2025-05-28T14:14:19Z)
• Centaur: Robust End-to-End Autonomous Driving with Test-Time Training [84.78837437133234]
  We propose Centaur, which updates a planner's behavior via test-time training without relying on hand-engineered rules or cost functions.<n>We develop a novel uncertainty measure, called Cluster Entropy, which is simple, interpretable, and compatible with state-of-the-art planning algorithms.
  arXiv  Detail & Related papers  (2025-03-14T17:59:41Z)
• Bridging Discrete and Backpropagation: Straight-Through and Beyond [62.46558842476455]
  We propose a novel approach to approximate the gradient of parameters involved in generating discrete latent variables. We propose ReinMax, which achieves second-order accuracy by integrating Heun's method, a second-order numerical method for solving ODEs.
  arXiv  Detail & Related papers  (2023-04-17T20:59:49Z)
• Efficient semidefinite bounds for multi-label discrete graphical models [6.226454551201676]
  One of the main queries on such models is to identify the SDPWCSP Function on Cost of a Posteri (MAP) Networks. We consider a traditional dualized constraint approach and a dedicated dedicated SDP/Monteiro style method based on row-by-row updates.
  arXiv  Detail & Related papers  (2021-11-24T13:38:34Z)
• Zeroth-order non-convex learning via hierarchical dual averaging [26.023679256204737]
  We propose a hierarchical version of dual dynamic averaging for zeroth-order online non- norm optimization. We derive tight bounds for both the learners static dynamic regret - i.e., the best policy in hindsight the play.
  arXiv  Detail & Related papers  (2021-09-13T09:59:06Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.