Hiroki Naganuma

Which Geometry on Which Layer? Measured Sign-Spectral Assignment in Mixed-Optimizer Training

Laura Gomezjurado

Mahdi Ghaznavi

Ganesh Talluri

Atsushi Nitanda

Seng Pei Liew

Ryuichiro Hataya

Modern neural-network training uses mixed recipes, assigning AdamW-like or sign-based updates to embedding and output layers and spectral up… (see more)dates to hidden matrices. These assignments are usually fixed by layer type, without a clear criterion for selecting the appropriate geometry. We turn sign-versus-spectral assignment into a measurable, layerwise decision problem. For ideal memoryless steps, we derive a finite-batch ratio of conservative descent-lemma lower bounds for spectral and sign updates. Its closed form separates signal efficiency from a geometry-dependent global curvature factor. In the positive-signal regime, the same smoothness assumption gives two-sided bounds that determine which step has the larger conditional expected decrease when the ratio crosses explicit noise-dependent thresholds. We then derive a local directional quadratic score from gradient alignments and Hessian-vector products and use it in Generalized-Scion (G-Scion) to control layerwise routing. In a controlled input-rank experiment, including directional curvature matches the measured same-iterate ordering in all seven aggregate settings, whereas the curvature-free proxy selects the opposite geometry. On 31M GPT-2, the output score routes the output layer from spectral warmup to sign, and G-Scion improves validation loss over fixed Scion in memoryless and matched-momentum comparisons. ViT-B/16 results further show that a spectral edge assignment can outperform fixed Scion outside language modelling. These results suggest that optimizer geometry need not be hard-coded by layer type and are a first step toward understanding how it should be assigned.

2026-10-11

arXiv (preprint)

Generalization Measures under Controlled Covariate Shift: A Regime-Aware Benchmark

Sora Nakai

Youssef Fadhloun

Kacem Mathlouthi

Kotaro Yoshida

Ganesh Talluri

Predicting generalization from quantities available before target-test evaluation remains a central challenge in deep learning. The systemat… (see more)ic benchmark of Jiang et al. (2020) evaluated many generalization measures, but it focused on independent and identically distributed (IID) settings. We revisit this problem for image classifiers evaluated under controlled corruptions and perturbations. Our study uses CIFAR-10-C/P, where the label space and task remain fixed while the input images are degraded or perturbed. This setting also allows us to revisit the robustness concerns raised by Dziugaite et al. (2020), who showed that the apparent reliability of generalization measures can depend strongly on experimental conditions. Our experiments show that the usefulness of generalization measures is strongly regime-dependent. In our exploratory decision analysis across three CNN-style architectures, sharpness- and input-gradient-based measures are among the leading individual signals, whereas family results are close and architecture dependent. Optimization-based measures, Information Criteria, and Sharpness-based measures provide additional regime-dependent signals in correlation or local-reliability analyses. Together, these findings suggest that model selection should not rely only on measures favored by IID evaluation. Instead, within the evaluated CIFAR-10-C/P setting and architectures, generalization measures should be treated as regime-dependent ranking signals whose utility must be evaluated for the intended corruption or perturbation setting.

2026-08-24

Transactions on Machine Learning Research (accepted)

Orth-Dion: Eliminating Geometric Mismatch in Distributed Low-Rank Spectral Optimization

Tatsuhiro Nakamori

Laura Gomezjurado Gonzalez

Ganesh Talluri

Ansh Tiwari

Hideyuki Kawashima

Guillaume Rabusseau

Low-rank gradient compression reduces communication in distributed training by representing updates with rank-…

2026-05-28

CoLoRAI @ International Conference on Machine Learning (poster)

Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent

Shagun Gupta

Youssef Briki

Irina Rish

Parameswaran Raman

Hao-Jun Michael Shi

To maximize hardware utilization, modern machine learning systems typically employ large constant or manually tuned batch size schedules, re… (see more)lying on heuristics that are brittle and costly to tune. Existing adaptive strategies based on gradient noise scale (GNS) offer a principled alternative. However, their assumption of SGD's Euclidean geometry creates a fundamental mismatch with popular optimizers based on generalized norms, such as signSGD / Signum (

2026-04-29

International Conference on Machine Learning (accepted)

The Geometry of Spectral Gradient Descent: Layerwise Criteria for SignSGD vs SpecSGD

Laura Gomezjurado

Mahdi Ghaznavi

Optimization in deep learning has expanded beyond Euclidean methods to include entrywise sign updates (SignSGD) and spectral sign updates (S… (see more)pecGD/Muon). While both can be viewed as steepest descent under non-Euclidean geometries (

2026-03-01

GRaM @ International Conference on Learning Representations (poster)

Takeuchi's Information Criteria as Generalization Measures for DNNs Close to NTK Regime

Taiji Suzuki

Rio Yokota

Masahiro Nomura

Kohta Ishikawa

Ikuro Sato

Generalization measures have been studied extensively in the machine learning community to better characterize generalization gaps. However,… (see more) establishing a reliable generalization measure for statistically singular models such as deep neural networks (DNNs) is difficult due to their complex nature. This study focuses on Takeuchi's information criterion (TIC) to investigate the conditions under which this classical measure can effectively explain the generalization gaps of DNNs. Importantly, the developed theory indicates the applicability of TIC near the neural tangent kernel (NTK) regime. In a series of experiments, we trained more than 5,000 DNN models with 12 architectures, including large models (e.g., VGG-16), on four datasets, and estimated the corresponding TIC values to examine the relationship between the generalization gap and the TIC estimates. We applied several TIC approximation methods with feasible computational costs and assessed the accuracy trade-off. Our experimental results indicate that the estimated TIC values correlate well with the generalization gap under conditions close to the NTK regime. However, we show both theoretically and empirically that outside the NTK regime such correlation disappears. Finally, we demonstrate that TIC provides better trial pruning ability than existing methods for hyperparameter optimization.

2026-02-25

arXiv (preprint)

Revisiting Generalization Measures Beyond IID: An Empirical Study under Distributional Shift

Sora Nakai

Youssef Fadhloun

Kacem Mathlouthi

Kotaro Yoshida

Ganesh Talluri

Generalization remains a central yet unresolved challenge in deep learning, particularly the ability to predict a model's performance beyond… (see more) its training distribution using quantities available prior to test-time evaluation. Building on the large-scale study of Jiang et al. (2020). and concerns by Dziugaite et al. (2020). about instability across training configurations, we benchmark the robustness of generalization measures beyond IID regime. We train small-to-medium models over 10,000 hyperparameter configurations and evaluate more than 40 measures computable from the trained model and the available training data alone. We significantly broaden the experimental scope along multiple axes: (i) extending the evaluation beyond the standard IID setting to include benchmarking for robustness across diverse distribution shifts, (ii) evaluating multiple architectures and training recipes, and (iii) newly incorporating calibration- and information-criteria-based measures to assess their alignment with both IID and OOD generalization. We find that distribution shifts can substantially alter the predictive performance of many generalization measures, while a smaller subset remains comparatively stable across settings.

2026-02-01

Open MIND (preprint)

DisTaC: Conditioning Task Vectors via Distillation for Robust Model Merging

Kotaro Yoshida

Yuji Naraki

Takafumi Horie

Ryotaro Shimizu

Model merging has emerged as an efficient and flexible paradigm for multi-task learning, with numerous methods being proposed in recent year… (see more)s. However, these state-of-the-art techniques are typically evaluated on benchmark suites that are highly favorable to model merging, and their robustness in more realistic settings remains largely unexplored. In this work, we first investigate the vulnerabilities of model-merging methods and pinpoint the source-model characteristics that critically underlie them. Specifically, we identify two factors that are particularly harmful to the merging process: (1) disparities in task vector norms, and (2) the low confidence of the source models. To address this issue, we propose **DisTaC** (**Dis**tillation for **Ta**sk vector **C**onditioning), a novel method that pre-conditions these problematic task vectors before the merge. DisTaC leverages knowledge distillation to adjust a task vector's norm and increase source-model confidence while preserving its essential task-specific knowledge. Our extensive experiments demonstrate that by pre-conditioning task vectors with DisTaC, state-of-the-art merging techniques can successfully integrate models that exhibit these harmful traits, where they would otherwise fail, and achieve significant performance gains.

2025-12-31

International Conference on Learning Representations (Accept (Poster))

Pseudo-Asynchronous Local SGD: Robust and Efficient Data-Parallel Training

Xinzhi Zhang

Man-Chung Yue

Russell J. Hewett

Philipp Andre Witte

Yin Tat Lee

2025-09-14

TMLR (accepted)

An Empirical Study of Pre-trained Model Selection for Out-of-Distribution Generalization and Calibration

Ryuichiro Hataya

Kotaro Yoshida

2025-05-12

TMLR (accepted)

No Wrong Turns: The Simple Geometry Of Neural Networks Optimization Paths

Charles Guille-escuret

Kilian Fatras

Understanding the optimization dynamics of neural networks is necessary for closing the gap between theory and practice. Stochastic first-or… (see more)der optimization algorithms are known to efficiently locate favorable minima in deep neural networks. This efficiency, however, contrasts with the non-convex and seemingly complex structure of neural loss landscapes. In this study, we delve into the fundamental geometric properties of sampled gradients along optimization paths. We focus on two key quantities, which appear in the restricted secant inequality and error bound. Both hold high significance for first-order optimization. Our analysis reveals that these quantities exhibit predictable, consistent behavior throughout training, despite the stochasticity induced by sampling minibatches. Our findings suggest that not only do optimization trajectories never encounter significant obstacles, but they also maintain stable dynamics during the majority of training. These observed properties are sufficiently expressive to theoretically guarantee linear convergence and prescribe learning rate schedules mirroring empirical practices. We conduct our experiments on image classification, semantic segmentation and language modeling across different batch sizes, network architectures, datasets, optimizers, and initialization seeds. We discuss the impact of each factor. Our work provides novel insights into the properties of neural network loss functions, and opens the door to theoretical frameworks more relevant to prevalent practice.

2024-07-22

International Conference on Machine Learning (Accept (Poster))