Machine Learning · Cross-listed from cs.CL
Evaluating LLM Personalization via Semantic Constraint Verification
Xuran Li, Guanqin Zhang, Imran Razzak et al. (+4 more)
Current evaluation paradigms for Large
Language Model (LLM) personalization rely heavily on brittle surface-matching metrics or computationally expensive LLM-as-a-judge protocols, both of which lack interpretability. To address these limitations, we introduce Natural
Language Inference Constraint Verification (NLICV), a scalable, semantically invariant framework that maps sentence meanings to truth-condition sets to verify personalization constraints via a Natural
Language Inference (NLI) model. Moving beyond binary scoring, NLICV categorizes LLM behaviors into four distinct modes: personalization, generalization, sycophancy, and failure. Extensive experiments demonstrate that NLICV aligns closely with human annotations while drastically reducing the latency and
token costs associated with LLM judges (up to 2100 inference speedup). Finally, through an ablation-based procedure, NLICV pinpoints the exact sentences driving the constraint verification, yielding faithful, understandable evidence for its evaluations.
Read the full abstract ▾
All authors Xuran Li, Guanqin Zhang, Imran Razzak, Hakim Hacid, Eleanna Kafeza, Hao Xue, Flora D. Salim
Machine Learning · Cross-listed from cs.CR
FEnc$^2$: Unifying Data Packing for Efficient Private Inference via Convolution and Architecture-Aware Fragment Encoding
Ran Ran, Zhaoting Gong, Nuo Xu et al. (+3 more)
Fully Homomorphic Encryption (FHE) enables
privacy-preserving machine learning but incurs extreme computational and
memory overhead. These costs come not only from expensive low-level primitives, including Number Theoretic Transform (NTT), rotation, and key-switching, but also from inefficient ciphertext packing at the application level. Existing packing strategies typically preserve either neighboring data elements or feature grouping, but not both, leading to wasted ciphertext slots, excessive rotations, and inflated ciphertext counts. We propose FEnc2, a unified and principled fragment-based encoding framework for CKKS-based private convolutional neural network inference. FEnc2 optimizes slot utilization, rotation complexity, and ciphertext density through two components: 1)Conv-aware Encoding, which analytically selects an optimal fragment size to decouple spatial dependencies and jointly minimize inner-outer rotations across layers, and 2)Arch-aware Ct Compression, which restores ciphertext density after feature- or channel-reduction layers. Together, these transformations reshape encrypted workload structure and reduce homomorphic operations by one to two orders of magnitude. With full
memory capacity utilized, i.e., at maximum batch size, FEnc2 achieves end-to-end latency speedups over the state-of-the-art Orion of up to 228.83x on GPU and 226.06x on CPU for LeNet on MNIST, and up to 4.55x on GPU and 9.43x on CPU for MobileNet on ImageNet. FEnc2 is hardware-agnostic yet architecturally transformative: by optimizing encrypted tensor layout before execution, it reduces ciphertext count and workload pressure on hardware, complementing primitive-level optimizations such as NTT and keyswitch accelerators. These results show that application-level data layout is a first-order architectural design dimension for encrypted inference and an important enabler for next-generation FHE systems.
Read the full abstract ▾
All authors Ran Ran, Zhaoting Gong, Nuo Xu, Yuanchao Xu, Fan Yao, Wujie Wen
Machine Learning
Simulation-Augmented Multi-Step Split Conformal Prediction for Aggregated Forecasts
Andro Sabashvili
We study uncertainty quantification for aggregated
forecasting tasks such as annual totals and year-over-year growth rates. We propose SA-MSCP, a simulation-augmented multi-step split conformal method that generates future paths from cross-validated residuals using a block bootstrap and constructs prediction intervals from empirical
quantiles. Experiments show that SA-MSCP improves empirical coverage over a simulated-path baseline for aggregated and growth-rate targets. Our results demonstrate that simulation-enhanced conformal calibration is an effective and general framework for uncertainty quantification in aggregated time-series
forecasting.
Read the full abstract ▾
Machine Learning
Communication-Efficient Verifiable Attention for LLM Inference
Ziqun Chen, Ming Wu, Michael Heinrich et al. (+4 more)
Computation integrity of remote large
language model (LLM) serving can be questionable. For conventional deep neural networks (DNNs), the existing TEE-shielded DNN partitioning (TSDP) approach uses Trusted Execution Environment (TEE) to compute non-linear components and verify the integrity of linear components offloaded to an untrusted GPU. However, directly applying TSDP to
Transformer-based LLMs incurs significant TEE computation and TEE-GPU communication overhead. This paper presents Communication-efficient TEE-GPU
Attention (VeriAttn) for accelerating verifiable LLM inference. VeriAttn offloads both linear and non-linear computations of
attention to the GPU, while TEE performs verification. Moreover, for prefill, VeriAttn uses a two-level pipeline to overlap data movement, TEE pre-/post-processing, and GPU computation. For decoding, when the key-value cache exceeds available GPU
memory, VeriAttn partitions
attention across TEE and GPU to reduce repeated key-value transfers. Evaluation on an Intel TDX platform shows that VeriAttn achieves 2.60-3.38$\times$ and 3.86-5.42$\times$ acceleration over TSDP for 6k-
token prompts and 10k-
token outputs during prefill and decoding, respectively.
Read the full abstract ▾
All authors Ziqun Chen, Ming Wu, Michael Heinrich, Jason Zeng, Huiying Lan, Tianwei Zhang, Rui Tan
Machine Learning · Cross-listed from cs.AI
Whose hotel does the AI recommend? An algorithm audit of reputation signals in LLM-assisted hotel selection
Mirza Samad Ahmed Baig, Syeda Anshrah Gillani and Asher Ali
Travelers increasingly ask large
language model (LLM) assistants which hotel to book, making these systems gatekeepers of property visibility – yet what moves their recommendations is undocumented. We conduct a pre-specified
algorithm audit using a randomized choice-based conjoint: across personas,
prompt templates, and twelve open-weight and proprietary models, assistants choose among five hotels whose guest rating, review volume and recency, management response, chain affiliation, price, eco-certification, and list position are independently randomized. We estimate the average marginal component effect of each signal on the probability of recommendation. Guest rating and price dominate (a top rating raises selection by 31.6 percentage points; a high price lowers it by 30.0), reproducing human valence-and-price primacy but over-weighting eco-certification and ignoring management response. List position – a content-free artifact – shifts recommendations causally, worth about \$12 per night. Stated reasons track revealed weights imperfectly. The findings ground generative engine optimization and the accountability of AI infomediaries in causal evidence.
Read the full abstract ▾
Machine Learning
Filtered ANN as a Phase Transition: When Selectivity-Estimation Error Causes Plan Regret
Madhulatha Mandarapu and Sandeep Kunkunuru
A filtered approximate-nearest-neighbor (ANN) query returns the k nearest vectors among those satisfying an attribute predicate P of selectivity s. The best execution strategy – pre-filter, post-filter, or in-filter – changes with s, so a system must estimate s and choose. We model this as an argmax over a landscape with phases (regions where each strategy wins) separated by boundaries, and show that selectivity-estimation error produces plan regret – recall lost versus the oracle strategy – only in the critical regions around those boundaries. The regret is a wedge of log-width equal to the multiplicative estimation error epsilon and height equal to the local cliff |V'(s*)| epsilon; the flip-margin 1/|V'(s*)| is the condition number of a sibling cardinality-estimation study reappearing as the local boundary theory. The two phase boundaries follow from independent mathematics: order statistics place the post-filter cliff at s ~ k/K, and site percolation places the in-filter cliff at s_c ~ 0.83/M for graph degree M (corpus-size independent). Criticality exists only under a constrained budget B < sqrt(k n). Under pre-registered decision rules we confirm, on synthetic sweeps and real SIFT1M, that regret concentrates ~290x at the boundary and that the regret curves obey a finite-size scaling collapse onto one universal wedge across two decades of corpus size. A real approximate index does not mis-locate the boundary, but a biased cost model opens a persistent miscalibration band that estimation-error robustness cannot fix. The contribution is a characterization, not a new index. Code and the full pre-registration are public.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.AI
Medical Heuristic Learning: An LLM-Driven Framework for Interpretable and Auditable Clinical Decision Rules
Wei Xu, Ke Yang, Gang Luo et al. (+4 more)
Predictive modeling for clinical decision support requires both strong predictive performance and transparent, auditable, and human-reviewable decision logic. Although deep learning and tree-based ensemble methods can achieve high accuracy, their black-box nature remains a major obstacle to trustworthy clinical deployment. Moreover, clinical prediction often operates under practical constraints, including limited sample sizes, severe class imbalance, and feature
evolution arising from changes in diagnostic criteria or clinical documentation practices. We propose Medical Heuristic Learning (MHL), a constrained paradigm for LLM-assisted rule learning. Rather than relying on updates to implicit model weights, MHL integrates statistical probes, medical knowledge probes, initial rule synthesis, and iterative rule optimization to construct an executable rule-based expert system. The resulting rule system is expressed entirely using the native logical and control-flow constructs of a programming
language. Valid rule versions are recorded and retained along the search trajectory, making the decision logic explicit, interpretable, and auditable. MHL also supports continual learning by using previously validated rules as a starting point and iteratively revising them in response to updated feature information under data drift or feature
evolution. MHL is not tied to any specific programming
language. Comprehensive experiments on medical
datasets show that MHL achieves predictive performance comparable to that of state-of-the-art methods, performs favorably in small-sample and highly imbalanced settings, and supports the transfer and adaptive revision of validated rules under feature
evolution. Overall, these findings suggest that non-gradient-based heuristic systems offer an approach to balancing predictive performance and transparency in clinical decision support.
Read the full abstract ▾
All authors Wei Xu, Ke Yang, Gang Luo, Keli Zheng, Lingyan Hu, Jing Wang, Kefeng Li
Machine Learning · Cross-listed from cs.CV
Differentiable Packing of Irregular 3D Objects with Adaptive Container Estimation
Palak Gupta and Shanmuganathan Raman
Most existing approaches either fix the container in advance or optimize only a single container dimension through an outer search loop, leaving the remaining dimensions as a manual tuning problem. We present a differentiable packing framework that jointly optimizes all 6N object pose parameters and all three container side lengths inside a single gradient-based loop. The formulation combines six physics-inspired, differentiable loss terms computed directly on triangle meshes through axis-aligned bounding-box proxies. An adaptive squeezing mechanism periodically tightens the container whenever the overlap loss falls below a pair-count-scaled threshold, producing a large initial drop in container volume, followed by small refinements. All pairwise computations are written in tensor-broadcasting form, giving a 3.4 to 54 times speedup over a reference loop-based implementation. The pipeline is implemented in Python and PyTorch, with no physics engine, FFT library, or convex decomposition. On multiple object categories, the method produces containers that are 11 to 32 percent smaller than time-matched DBLF and simulated-annealing baselines at N =100, while running in under 4 minutes per instance on a single consumer GPU.
Read the full abstract ▾
Machine Learning
Diffusion Offline Reinforcement Learning for Fair and Energy-Efficient UAV-Assisted Wireless Networks
Eslam Eldeeb and Hirley Alves
The integration of generative artificial intelligence with wireless communication and signal processing systems has opened new avenues for intelligent, data-driven decision-making in future 6G networks. This work proposes a
diffusion soft actor-critic (
Diffusion-SAC) approach that leverages offline reinforcement learning (RL) enhanced by denoising
diffusion probabilistic models (DDPMs) to optimize trajectory and scheduling control in unmanned aerial vehicle (
UAV) networks. While offline RL methods, such as conservative Q-learning (CQL), can learn from static
datasets, they often struggle to generalize in low-data or dynamic conditions. To address this, we combine the robustness of CQL with the generative power of
diffusion models, enabling expressive and signal-aware policy learning that generalizes beyond behavior policies. Applied to a
UAV-assisted wireless network, the proposed framework minimizes transmission energy and improves
fairness among devices. Simulations show that
Diffusion-SAC outperforms standard offline RL baselines, achieving more stable convergence and higher rewards even with limited
datasets. The method enhances data efficiency, reduces energy consumption, and increases throughput by more than 35 % compared to existing
algorithms, demonstrating its potential for robust policy learning in next-generation wireless control systems.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.IR
RL-Index: Reinforcement Learning for Retrieval Index Reasoning
Yongjia Lei, Nedim Lipka, Zhisheng Qi et al. (+7 more)
Retrieving external knowledge is crucial for real-world tasks but remains difficult when queries and relevant knowledge are linked by implicit
reasoning (e.g., shared
theorems or coding logic). Existing methods rely mainly on query-side
reasoning, leading to high online latency and underutilizing the
reasoning semantics within the knowledge corpus. In this paper, we propose $\textbf{RL-Index}$, an indexing framework that formulates
retrieval index
reasoning as a reinforcement learning problem. Instead of performing
reasoning at query time, RL-Index shifts
reasoning to the indexing stage by augmenting documents with LLM-generated rationales that explicitly encode the latent query-knowledge relationship. To optimize the quality of these rationales, we employ Group Relative Policy Optimization (GRPO) and use
retrieval similarity as a proxy reward signal, enabling direct optimization of indexing decisions for
retrieval effectiveness. Extensive experiments on the BRIGHT
benchmark demonstrate that RL-Index consistently improves both
retrieval and downstream question-answering performance, while significantly reducing online inference latency. Moreover, the learned rationale augmentation generalizes across diverse retrievers and generators, highlighting its robustness as a plug-and-play indexing strategy across different
retrieval systems.
Read the full abstract ▾
All authors Yongjia Lei, Nedim Lipka, Zhisheng Qi, Utkarsh Sahu, Yuchen Zhuang, Wenqi Shi, Koustava Goswami, Franck Dernoncourt, Ryan A. Rossi, Yu Wang
Machine Learning
QK-Normed MLA: QK normalization without full key caching
Yizhou Han, Yao Zhao, Jun Zhou et al. (+2 more)
Query-key (QK) normalization stabilizes
attention by controlling the scale of queries and keys before the dot product, but is not immediately compatible with Multi-head Latent
Attention (MLA). MLA achieves efficient decoding by caching low-dimensional latent states instead of full keys, whereas post-projection QK RMSNorm appears to require the fully projected key for every cached
token. We show this apparent incompatibility is an implementation artifact, not an architectural constraint. RMSNorm decomposes into a static affine weight and a dynamic scalar RMS statistic. The static key-side weight can be absorbed into the MLA query-side projection; the dynamic key statistic reduces to one inverse-RMS scalar per
token and KV group. The resulting formulation is exactly equivalent to explicit post-projection QK RMSNorm in exact arithmetic and preserves MLA's latent decode path. In our 400M runs trained for up to 100B
tokens, QK-Normed MLA achieves lower training loss and better downstream accuracy than QK clipping, while H800 decode
benchmarks show less than 2% latency overhead up to 256k context. These results make QK normalization a practical stabilization option for MLA models without requiring full-key caching.
Read the full abstract ▾
All authors Yizhou Han, Yao Zhao, Jun Zhou, Longfei Li, Ruoyu Sun
Machine Learning
pFedUL: Layer-Aware Federated Unlearning for Personalized Federated Learning
Zhuodong Liu, Xiangyu Li and Zhihao Zhang
Federated unlearning (FU) enables the removal of specific data contributions from
federated learning (FL) models to comply with regulations such as the General Data Protection Regulation (GDPR). However, most existing FU methods are designed for the FedAvg paradigm, where all clients share a single global model. In practice, personalized
federated learning (pFL) methods such as FedPer, FedRep, Ditto, and FedBN have become widely adopted due to their superior handling of non-IID data. These methods decompose the model into shared global layers and client-specific personalized layers, fundamentally altering the semantics of unlearning, yet this setting has received little
attention. We formalize FU under the pFL paradigm, identifying a tension between unlearning completeness on shared layers and personalization preservation for remaining clients. We then propose pFedUL, a layer-aware selective unlearning framework comprising three components: (1) gradient-based layer-wise contribution attribution that separately quantifies the target client's influence on shared and personalized parameters, (2) adaptive selective unlearning that applies differentiated forgetting strategies across layer types, and (3) a lightweight recalibration protocol enabling remaining clients to restore personalization with minimal overhead. We further introduce two new metrics, Personalization Preservation Score (PPS) and Cross-client
Fairness Index (CFI), to evaluate pFL-specific unlearning quality. Experiments on CIFAR-10, CIFAR-100, and FEMNIST under varying non-IID settings indicate that pFedUL achieves unlearning effectiveness comparable to full retraining while maintaining an average of 97.3% personalized accuracy for remaining clients. Compared with six state-of-the-art FU methods adapted to the pFL setting, pFedUL consistently achieves superior personalization preservation.
Read the full abstract ▾
Machine Learning
One-Step Generalization Ratio Guided Optimization for Domain Generalization
Sumin Cho, Dongwon Kim and Kwangsu Kim
Domain Generalization (DG) aims to train models that generalize to unseen target domains but often overfit to domain-specific features, known as undesired correlations. Gradient-based DG methods typically guide gradients in a dominant direction but often inadvertently reinforce spurious correlations. Recent work has employed dropout to regularize overconfident parameters, but has not explicitly adjusted gradient
alignment or ensured balanced parameter updates. We propose GENIE (Generalization-ENhancing Iterative Equalizer), a novel optimizer that leverages the One-Step Generalization Ratio (OSGR) to quantify each parameter's contribution to loss reduction and assess gradient
alignment. By dynamically equalizing OSGR via a preconditioning factor, GENIE prevents a small subset of parameters from dominating optimization, thereby promoting domain-invariant feature learning. Theoretically, GENIE balances convergence contribution and gradient
alignment among parameters, achieving higher OSGR while retaining SGD's convergence rate. Empirically, it outperforms existing optimizers and enhances performance when integrated with various DG and single-DG methods.
Read the full abstract ▾
Machine Learning
An affordable hardware-aware neural architecture search for deploying convolutional neural networks on ultra-low-power computing platforms
Andrea Mattia Garavagno, Edoardo Ragusa, Antonio Frisoli et al. (+1 more)
Hardware-aware neural architecture search (HW-NAS) allows the integration of Convolutional Neural Networks (CNNs) in microcontrollers devices by automatically designing neural architectures that can fit prearranged hardware constraints. However, state-of-the-art HW-NAS target high-performance microcontrollers, whose power consumption does not meet sensing nodes requirements. This work presents a HW-NAS generating tiny CNNs that can run on ultra-low-power microcontrollers, featuring a lightweight search procedure enabling its execution even on embedded devices. Empirical results on three well-known
benchmarks for tiny computer
vision proved that the proposed HW-NAS was able to generate tiny CNNs while preserving state-of-the-art classification accuracy.
Read the full abstract ▾
All authors Andrea Mattia Garavagno, Edoardo Ragusa, Antonio Frisoli, Paolo Gastaldo
Machine Learning
FlowMPC: Improving Flow Matching policies with World Models
Chandon Hamel
Flow Matching (FM) is a powerful approach for behavior cloning in
multimodal action spaces [Jiang et al., 2025], but because it is not trained to directly maximize expected return, there is still room to improve how FM policies act at test time. This work investigates whether a learned world model can improve FM policies by enabling Model Predictive Path Integral (MPPI) planning over candidate action sequences proposed by the policy. Building on TD-MPC2 [Hansen et al., 2024], I introduce FlowMPC, a framework that combines an imitation-learned FM policy with a learned world model for test-time planning in ManiSkill manipulation tasks [Tao et al., 2025]. Across PickCube and PickSingleYCB, adding the world model improved performance over the FM policy alone, with especially clear gains in end-of-episode success. These results suggest that world-model-based planning can effectively complement flow-based imitation policies without modifying the FM training objective.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
HiMPO: Hindsight-Informed Memory Policy Optimization for Less-Entangled Credit in Long-Horizon Agents
Jiangze Yan, Yi Shen, Wenjing Zhang et al. (+6 more)
Long-horizon
agents compress interaction histories into
memory, but final outcomes entangle
memory-writing credit with tool failures and
reasoning errors. We propose HiMPO, a Hindsight-Informed
Memory Policy Optimization framework that supplements trajectory-level reinforcement learning with
memory-specific credit. For each update, HiMPO compares the target answer's log-likelihood under the previous and updated
memories while holding the remaining pre-write context fixed. A soft gate based on group-relative hindsight relevance then weights the local utility advantage. Trajectory-level credit applies to all generated
tokens;
memory tokens additionally receive the auxiliary advantage. Across judge-based open-domain tasks and objective compressive-
memory QA, HiMPO improves task performance while preserving compressed-context efficiency. On Local Wiki Search, it improves average F1 and exact match over MemPO by 2.6 and 2.2 percentage points, respectively. Controlled interventions and live
memory-reversion studies provide evidence of reduced blame leakage and
alignment between
memory credit and behavioral effects. With the policy fixed, the credit
estimator remains robust to the tested forms of target noise.
Read the full abstract ▾
All authors Jiangze Yan, Yi Shen, Wenjing Zhang, Jieyun Huang, Ping Chen, Zhaoxiang Liu, Ning Wang, Kai Wang, Shiguo Lian
Machine Learning · Cross-listed from stat.ML
Generative Modeling on Metric Graphs via Neural Optimal Transport
Alessandro Micheli, Yueqi Cao, Anthea Monod et al. (+1 more)
We introduce, to our knowledge, the first deep generative modeling framework for probability distributions continuously supported on compact metric graphs. Given source and target measures on a metric graph, our method embeds the graph into a smooth ambient space, solves an entropic Kantorovich problem via a neural semidual parameterization, and projects generated samples back onto the original graph. We study two embedded geometries: an extrinsic Euclidean realization and the intrinsic tropical Abel–Jacobi
embedding into the Jacobian
torus. In both cases, the resulting generator is graph-supported by construction. We prove that, in the joint limit of increasing neural expressivity, the learned generator converges weakly to a valid transport coupling between the original graph measures. Empirically, across a range of geometrically distinct graphs, our method matches or improves upon heuristic transport baselines based on discrete graph OT, while scaling more favorably. Finally, we demonstrate scalability on real-world urban mobility data by training our model on one million Uber pickup locations in Manhattan, New York City.
Read the full abstract ▾
All authors Alessandro Micheli, Yueqi Cao, Anthea Monod, Samir Bhatt
Machine Learning · Cross-listed from cs.CV
Contrastive Learning for Seismic Horizon Tracking with Domain-Specific Priors
Alexandre Thouvenot, Lionel Boillot and Vincent Gripon
Unsupervised 3D seismic horizon tracking faces a key limitation: signal-based propagators provide accurate trace-level
alignment but often fail near faults, whereas texture-driven deep models are more robust to discontinuities, typically at the cost of labeled data requirements and reduced trace-level precision. We propose a self-supervised fusion of both paradigms in which signal-derived local horizon correspondences act as domain-specific priors to train a texture-based deep learning model. Specifically, we estimate reliable trace-to-trace flows from reflector slopes and use them to form positive pairs in a contrastive objective, while restricting training to high-confidence neighborhoods, optionally augmented with a fault mask. The objective is not to infer ambiguous correspondences close to discontinuities, but to preserve horizon identity across them. As a result, the network learns voxel-wise
embeddings that preserve local signal continuity while enabling horizon propagation beyond discontinuities through similarity search. Experiments on the public F3
dataset and a faulted synthetic
dataset achieve lower mean absolute error (MAE) than unsupervised baselines and competitive performance against a semi-supervised method using a single labeled slice.
Read the full abstract ▾
Machine Learning
Variance Reduction for Non-Log-Concave Sampling with Applications to Inverse Problems
M. Berk Sahin, Ahmet Ege Tanriverdi, Behzad Sharif et al. (+1 more)
Sampling from high-dimensional, non-log-concave distributions with unnormalized densities is a fundamental challenge in machine learning, particularly when the exact gradient of the potential is unavailable and must be approximated via
stochastic gradients that exhibit high variance under a fixed budget of gradient computations per iteration. Although variance reduction techniques such as SGD with momentum, STORM, and PAGE have demonstrated improved convergence properties in non-convex optimization, their implications for sampling from non-log-concave distributions remain largely unexplored. In this work, we develop the first unified analysis of these
estimators for sampling from non-log-concave distributions. We establish improved non-asymptotic convergence rates in $\varepsilon$-relative Fisher information and, under a Poincaré inequality assumption, in squared total variation distance, and further prove weak convergence to the target distribution. We extend our analysis to solving inverse problems with score-based generative priors. We empirically validate our theory and demonstrate that, under a fixed gradient computations per iteration, variance-reduction techniques consistently improve sample quality in two standard imaging applications.
Read the full abstract ▾
All authors M. Berk Sahin, Ahmet Ege Tanriverdi, Behzad Sharif, Abolfazl Hashemi
Machine Learning · Cross-listed from cs.CV
KeepLoRA++: Continual Learning with Layer-Scaled Residual Gradient Adaptation
Mao-Lin Luo, Yi-Lin Zhang, Zi-Hao Zhou et al. (+5 more)
Continual learning for pre-trained
vision-
language models requires balancing three competing objectives: retaining pre-trained knowledge, preserving knowledge from a sequence of learned tasks, and maintaining the plasticity to acquire new knowledge. This paper presents KeepLoRA++, balancing these objectives through a unified dual-dimensional knowledge retention mechanism. We analyze knowledge distribution of
Transformer architecture from both inter-layer and intra-layer perspectives. The inter-layer perspective examines how retention is distributed across layers, while the intra-layer perspective focuses on the parameter space within each layer. Our analysis reveals a structural property: general transferable knowledge is mainly encoded in the shallow layers and the principal subspace of the parameters, while task-specific adaptations are localized in the deep layers and the residual subspace. Motivated by this insight, KeepLoRA++ introduces a layer-scaled residual gradient adaptation method. New tasks are learned by restricting LoRA parameter updates to the residual subspace, combined with a shallow-to-deep layer scaling, to prevent interference with previously acquired capabilities. Specifically, the gradient of a new task is projected onto a subspace orthogonal to both the principal subspace of the pre-trained model and the dominant directions of previous task features, while simultaneously assigning smaller update magnitudes to shallow layers and larger ones to deeper layers. Our theoretical analysis and empirical evaluations confirm that KeepLoRA++ successfully balances these three competing objectives, consistently outperforming representative baselines across image classification, visual question answering, and
video understanding tasks.
Read the full abstract ▾
All authors Mao-Lin Luo, Yi-Lin Zhang, Zi-Hao Zhou, Yankun Hong, Xialiang Tong, Mingxuan Yuan, Tong Wei, Min-Ling Zhang
Machine Learning
Demystifying Training-Time Augmentation for Data-Constrained Language Model Pretraining
Michael K. Chen, Xikun Zhang, Fan Bai et al. (+2 more)
As AI labs approach a data ceiling where compute capacity outpaces the rate of new high-quality text generation,
language model pretraining is shifting toward a data-constrained, compute-abundant regime that demands productive multi-epoch training on fixed corpora. Standard autoregressive (AR) pretraining overfits severely in this setting, reaching its optimum early and then continuously deteriorating. We investigate training-time data augmentation as a regularizer to mitigate this overfitting and enable productive training for hundreds of epochs on the same data. We introduce three orthogonal categories of augmentation for AR pretraining:
token-level noise (masking, random replacement), sequence permutations (right-to-left prediction, Fill-in-the-Middle), and target offset prediction ($x_{t+i}$ for $i > 1$). Through systematic ablations, we find that individual augmentations delay overfitting and lower validation loss relative to the baseline, with random
token replacement achieving the best minimum loss among individual methods. Combining augmentation categories further lowers the minimum validation loss. Our experiments demonstrate that data augmentations mitigate AR pretraining's data inefficiency and offer a promising solution to the data-constrained regime~\footnote{All code and data are available at https://github.com/ michaelchen-lab/ data-augmentations-for-pretraining.
Read the full abstract ▾
All authors Michael K. Chen, Xikun Zhang, Fan Bai, Zhengding Hu, Zhen Wang
Machine Learning
LiFT: Local Search via Linear Programming for Overfitting-Controlled Transformers
Abhishek Shukla, Anikeit Khanna, Ankur Sinha et al. (+1 more)
This paper proposes a Linear Programming (LP)-based local search framework for fine-tuning pretrained
transformer models with explicit control against overfitting. The approach formulates
transformer fine-tuning as a bilevel optimization-based regularization problem, in which model parameters and regularization hyperparameters are jointly updated. Information collected during initial warm-up iterations, including validation gradients and training Hessian information, is used to construct a local descent direction by solving an LP that minimizes a scaled directional derivative while preserving training optimality. This validation-aware descent direction enables focused local updates of both parameters and regularization hyperparameters, reducing overfitting without requiring repeated full retraining cycles. The resulting method, termed Linear Programming-based Fine-Tuning (LiFT) for
transformers, differs from conventional fine-tuning by systematically identifying task-specific updates rather than relying on heuristic or grid-based hyperparameter selection. Experiments on GPT-2 Small fine-tuned on WikiText-2 demonstrate that LiFT enables effective adaptation through selective tuning of
transformer blocks and regularization parameters, yielding consistent improvements in test perplexity across multiple layer configurations and regularization settings, with particularly pronounced gains in overfitting-prone scenarios. Beyond empirical performance, LiFT establishes a principled connection between
transformer fine-tuning, bilevel optimization, local search, and regularization theory.
Read the full abstract ▾
All authors Abhishek Shukla, Anikeit Khanna, Ankur Sinha, Faiz Hamid
Machine Learning
Rapid Poison: Practical Poisoning Attacks Against the Rapid Response Framework
David Huang, Jaewon Chang, Avidan Shah et al. (+2 more)
The Rapid Response (RR) framework, deployed in production systems, including Anthropic's ASL-3 safeguards, continuously improves jailbreak-detection classifiers. When new jailbreaks emerge that bypass these classifiers, Rapid Response generates synthetic variants for training, helping the model generalize from the new attacks and quickly adapt. We reveal that
prompt injection can infiltrate this pipeline to deliver poisoned samples into the classifier's training set, enabling two attack objectives: (I) targeted poisoning attacks that create false positives on harmless samples by categorizing them as a jailbreak, with a specific desired feature (e.g., certain formatting, subject, or keyword), (II) concept-based
backdoor attacks that induce false negatives on jailbreak inputs, generalizing even to jailbreaks from attack strategies the defender explicitly trained against, when the
backdoor trigger is present. Importantly, our threat model restricts adversaries to modifying only jailbreak samples (not benign data or labels), a constraint unexplored by prior work that makes the second objective particularly challenging. We address this with Omission Attack, which exploits a new phenomenon: when training on concept-absent unsafe samples, the classifier misassociates that concept's presence with the safe label. Both attacks cause substantial and in some cases near-complete label flipping at only a 1% poisoning rate, achieving up to 100% false positive rates and up to 96% false negative rates.
Read the full abstract ▾
All authors David Huang, Jaewon Chang, Avidan Shah, Prateek Mittal, Chawin Sitawarin
Machine Learning · Cross-listed from cs.CL
Creative Collision: Directorial Persona Steering and Competition in Large Language Models
Subramanyam Sahoo and Justin Shenk
Activation steering has emerged as a powerful tool for shaping the behaviour of large
language models at inference time, yet most prior work injects a single semantic direction into the residual stream. We study the richer setting in which two semantically opposing steering vectors are superimposed – a regime we call Creative Collision. Concretely, we construct directorial persona vectors for Steven Spielberg (optimistic, redemptive moral valence) and Martin Scorsese (dark, morally ambiguous) via mean-difference activation contrast on curated screenplay-derived corpora, then interpolate between them with a scalar mixing parameter $\alpha \in [0,1]$ and a steering coefficient $\lambda$. Across five evaluation axes – moral valence, generation coherence, surface style, directional dominance, and vector geometry – three principal findings emerge: (i)~Spielberg's representational signature exhibits robust directional dominance, suppressing Scorsese's moral influence across almost the entire interpolation range; (ii)~intermediate collision points paradoxically improve generation coherence relative to pure single-director steering at high $\lambda$; and (iii)~both personas localise maximally to layer~28 of a 40-layer decoder-only
transformer, revealing a shared moral-tone substrate. These results illuminate the geometry of competing semantic directions in
transformer residual streams and have direct implications for controllable creative generation and value-aligned narrative synthesis.
Read the full abstract ▾
Machine Learning
Evolutionary Bilevel Reward Shaping for Generalization in Reinforcement Learning
Ekasit Usaratniwart, Xilin Gao, Marc Ong et al. (+1 more)
Reinforcement learning (RL) often suffers from performance degradation when deployed in environments that differ from those encountered during training. Existing techniques such as domain randomization (DR) mitigate this, but require access to diverse training environments and full trajectory observability, assumptions that fail in
privacy-preserving or restricted scenarios where only scalar performance metrics are available. We propose Generalization via Evolutionary Reward Shaping (GERS), a bilevel optimization approach to improve generalization on unseen test environments using only scalar feedback from validation environments. At the lower level, an RL
agent guided via a reward function shaped by the upper level learns a policy on a limited set of training environments with accessible trajectory data; at the upper level, CMA-ES optimizes the reward shaping parameters to maximize the cumulative unshaped reward on separate validation environments for which trajectory access is unavailable. Results on continuous control tasks indicate that GERS outperforms the standard RL baseline on unseen test environments. GERS performance is comparable to DR, despite DR treating the combined set of training and validation environments of GERS as a single training set that requires trajectory access, whereas GERS cannot access validation trajectories. These results confirm that GERS effectively enhances generalization under restricted data access constraints.
Read the full abstract ▾
All authors Ekasit Usaratniwart, Xilin Gao, Marc Ong, Youhei Akimoto
Machine Learning
From Tokens to Regions: CUDA-Sensitive Instruction Tuning for GPU Kernel Generation
Wentao Chen, Jiace Zhu, Xing Zhe Chai et al. (+4 more)
High-performance CUDA kernels are essential for scalable AI systems, while Large
Language Models (LLMs) still struggle to generate correct kernels due to strict and implicit execution constraints. Existing LLM-based approaches either rely on costly agentic or reinforcement-learning (RL) pipelines, or adopt supervised fine-tuning (SFT) objectives that fail to explicitly model CUDA sensitivity, namely code
tokens or regions tightly coupled with execution constraints. In this work, we investigate CUDA sensitivity from the perspective of
token confidence patterns, showing that CUDA sensitivity appears at both
token and region levels, where most CUDA-sensitive
tokens are predicted with high confidence, while a smaller low-confidence subset forms regions corresponding to execution-critical structures. These findings suggest that effective CUDA kernel generation should both leverage high-confidence CUDA-sensitive
tokens and preserve low-confidence CUDA-sensitive regions. Building on these insights, we propose \textbf{\underline{CU}DA-\underline{Se}nsitive Instruction \underline{T}uning (CuSeT)}, a low-cost post-training method within a simple SFT framework. CuSeT follows the principle of “from
tokens to regions” by combining adaptive
token-level masking with region-aware sample reweighting. Experiments show that CuSeT consistently improves functional correctness across multiple model families and scales, outperforming standard SFT and advanced SFT variants, while achieving competitive performance against frontier CUDA kernel generation models with substantially lower inference cost.
Read the full abstract ▾
All authors Wentao Chen, Jiace Zhu, Xing Zhe Chai, Zeng Qu, Qiaoling Xiao, Liucheng Duan, An Zou
Machine Learning
Prediction of Runtime Parameters of Parallel Chemistry Applications via Active and Generative Learning
Tanzila Tabassum, Omer Subasi, Ajay Panyala et al. (+5 more)
In this work, we develop two main Machine Learning based approaches to predict the runtime parameters of highly scalable parallel chemistry computations.These approaches employ active and generative learning together with the empirically determined gradient boosted regression tree models chosen among a rich suite of machine learning models. When evaluated on Coupled-Cluster with Singles and Doubles computations, our models achieve a mean absolute error percentage (MAPE) as low as 0.023 and a coefficient of determination as high as 99.9%. Furthermore, when combined with active learning to mitigate the lack of large amounts of training data, our models score a MAPE about 0.2 with 20-25% of the original
dataset.
Read the full abstract ▾
All authors Tanzila Tabassum, Omer Subasi, Ajay Panyala, Epiya Ebiapia, Gerald Baumgartner, Erdal Mutlu, P Sadayappan, Karol Kowalski
Machine Learning · Cross-listed from cs.AI
Latent Thought Flow: Efficient Latent Reasoning in Large Language Models
Xiandong Zou, Jing Huang, Jianshu Li et al. (+1 more)
Large
Language Models (LLMs) increasingly rely on intermediate
reasoning, yet explicit Chain-of-Thought (CoT) suffers from a linguistic space bottleneck: each thought must be decoded into
tokens, causing high inference overhead. Latent
reasoning moves deliberation into continuous space, but existing methods mostly learn deterministic or reward-maximizing paths, lacking a principled way to allocate probability across trajectories with different correctness and costs. We propose Latent Thought Flow (LTF), which models
reasoning as variable-length continuous trajectories and trains a sampler to match a reward-induced posterior over answer quality and computation cost. We instantiate this with a continuous GFlowNet using
stochastic latent transitions. To handle sparse answer supervision, we introduce an
Entropy-Weighted Subtrajectory Balance objective for intermediate rewards and a reference-prior regularizer to anchor exploration. Experiments under finetuning and transfer learning settings show that LTF outperforms explicit CoT and latent
reasoning baselines, improving accuracy by 9.5% while reducing
reasoning length by 27.2% on average compared with strong latent
reasoning baselines.
Read the full abstract ▾
All authors Xiandong Zou, Jing Huang, Jianshu Li, Pan Zhou
Machine Learning · Cross-listed from cs.CE
Sparse stochastic digital twins can outperform exact ones
Zongren Zou, Théo Bourdais, Ricardo Baptista et al. (+1 more)
Digital twins are often built under the assumption that greater model fidelity leads to better downstream performance. We show that this is not always the case: retaining weakly dependent or auxiliary degrees of freedom can increase the dimensionality of learning and control problems enough that a sparse
stochastic surrogate outperforms even an exact high-dimensional model. To identify such structure, we introduce a framework for dependence discovery directly at the level of the full conditional law. The method combines conditional generative modeling with Gaussian-process
distillation and kernel mode decomposition to quantify how strongly individual variables enter the conditional distribution, including through variability, tails, multimodality, and other features invisible to the conditional mean. The resulting dependence structure can be used both to construct sparse
stochastic digital twins from data and to reduce known high-dimensional models for downstream tasks. In a modified Lunar Lander problem, an eight-dimensional sparse surrogate yields better control performance than the exact 109-dimensional simulator despite lower transition fidelity. A
stochastic heat-equation example shows that the framework recovers sparse structure in a 102-dimensional data-driven setting and produces a more effective surrogate than one conditioned on all candidate inputs, while economic data demonstrate its use for identifying time-varying
memory structure. These results show that model fidelity and downstream utility are distinct objectives, and that dependence discovery provides a practical route to sparse
stochastic digital twins.
Read the full abstract ▾
All authors Zongren Zou, Théo Bourdais, Ricardo Baptista, Houman Owhadi
Machine Learning · Cross-listed from cs.CL
PACT: Privileged Trace Co-Training for Multi-Turn Tool-Use Agents
Zhenbang Du, Jun Luo, Zhiwei Zheng et al. (+8 more)
Multi-turn tool-use
agents must reason, call tools, and adapt to observations across several interaction turns. Post-training such
agents is challenging, as reinforcement learning often suffers from sparse rewards and weak credit assignment despite matching the
prompt-only inference setting, while supervised fine-tuning on expert traces provides dense process supervision but can over-constrain the model to fixed trajectories. To tackle this, we propose PACT, a Privileged trAce Co-Training framework for multi-turn tool-use
agents. The key idea is to use expert traces only as training-time optimization signals rather than rollout-time hints. PACT keeps rollout generation
prompt-only, then uses expert traces to guide optimization through two complementary signals: a trace-conditioned RL surrogate that evaluates
prompt-only rollouts under expert-trace context, and a component-aware SFT loss that provides annealed process supervision over
reasoning and tool-call components. To reduce over-reliance on the training-only trace context, PACT incorporates
prompt-only anchoring with standard
prompt-only RL updates. Experiments on FTRL, BFCL, and ToolHop show that PACT consistently outperforms strong SFT- and RL-based baselines. Further analyses support the complementarity of the two training objectives and the benefits of component-aware supervision and
prompt-only anchoring. These results highlight the effectiveness of privileged trace co-training for multi-turn tool-use learning. The project page is available at https://zhenbangdu.github.io/pact-project-page.
Read the full abstract ▾
All authors Zhenbang Du, Jun Luo, Zhiwei Zheng, Xiangchi Yuan, Kejing Xia, Dachuan Shi, Qirui Jin, Qijia He, Shaofeng Zou, Yingbin Liang, Wenke Lee