Eugene Belilovsky

pietro.cagnasso@mila.quebec

Xiao Huang

Maîtrise recherche - Concordia

Co-superviseur⋅e :

xiaolong.huang@mila.quebec

Paul Janson

Doctorat - Concordia

Co-superviseur⋅e :

paul.janson@mila.quebec

Gwen Legate

Doctorat - Concordia

Co-superviseur⋅e :

gwendolyne.legate@mila.quebec

Paria Mehrbod

Maîtrise recherche - Concordia

Co-superviseur⋅e :

paria.mehrbod@mila.quebec

Abhinav Moudgil

Doctorat - Concordia

abhinav.moudgil@mila.quebec

Adel Nabli

Doctorat - Concordia

adel.nabli@mila.quebec

Geraldin Nanfack

Postdoctorat - Concordia

Co-superviseur⋅e :

geraldin.nanfack@mila.quebec

Albert Orozco Camacho

Doctorat - Concordia

Co-superviseur⋅e :

Doctorat - Concordia

Co-superviseur⋅e :

vaibhav.singh@mila.quebec

Benjamin Therien

Doctorat - UdeM

Superviseur⋅e principal⋅e :

benjamin.therien@mila.quebec

Stagiaire de recherche - Concordia

junior-cedric.tonga@mila.quebec

Zhang Yanlei

Collaborateur·rice alumni - UdeM

Superviseur⋅e principal⋅e :

yanlei.zhang@mila.quebec

amirhossein.zamani@mila.quebec

Amirhossein Zamani

Doctorat - Concordia

Co-superviseur⋅e :

Blake Richards

Publications

Rethinking Prompt Optimization: Reinforcement, Diversification, and Migration in Blackbox LLMs

MohammadReza Davari

Utkarsh Garg

Weixin Cai

2025-07-13

ArXiv (prépublication)

arxiv.org

Circuit Discovery Helps To Detect LLM Jailbreaking

Despite extensive safety alignment, large language models (LLMs) remain vulnerable to jailbreak attacks that bypass safeguards to elicit har… (voir plus)mful content. While prior work attributes this vulnerability to safety training limitations, the internal mechanisms by which LLMs process adversarial prompts remain poorly understood. We present a mechanistic analysis of the jailbreaking behavior in a large-scale, safety-aligned LLM, focusing on LLaMA-2-7B-chat-hf. Leveraging edge attribution patching and subnetwork probing, we systematically identify computational circuits responsible for generating affirmative responses to jailbreak prompts. Ablating these circuits during the first token prediction can reduce attack success rates by up to 80\%, demonstrating its critical role in safety bypass. Our analysis uncovers key attention heads and MLP pathways that mediate adversarial prompt exploitation, revealing how important tokens propagate through these components to override safety constraints. These findings advance the understanding of adversarial vulnerabilities in aligned LLMs and pave the way for targeted, interpretable defenses mechanisms based on mechanistic interpretability.

2025-06-29

ICML.cc/2025/Workshop/R2-FM (poster)

Less is More: Undertraining Experts Improves Model Upcycling

Stefan Horoi

Gintare Karolina Dziugaite

Modern deep learning is increasingly characterized by the use of open-weight foundation models that can be fine-tuned on specialized dataset… (voir plus)s. This has led to a proliferation of expert models and adapters, often shared via platforms like HuggingFace and AdapterHub. To leverage these resources, numerous model upcycling methods have emerged, enabling the reuse of fine-tuned models in multi-task systems. A natural pipeline has thus formed to harness the benefits of transfer learning and amortize sunk training costs: models are pre-trained on general data, fine-tuned on specific tasks, and then upcycled into more general-purpose systems. A prevailing assumption is that improvements at one stage of this pipeline propagate downstream, leading to gains at subsequent steps. In this work, we challenge that assumption by examining how expert fine-tuning affects model upcycling. We show that long fine-tuning of experts that optimizes for their individual performance leads to degraded merging performance, both for fully fine-tuned and LoRA-adapted models, and to worse downstream results when LoRA adapters are upcycled into MoE layers. We trace this degradation to the memorization of a small set of difficult examples that dominate late fine-tuning steps and are subsequently forgotten during merging. Finally, we demonstrate that a task-dependent aggressive early stopping strategy can significantly improve upcycling performance.

2025-06-16

arXiv (prépublication)

arxiv.org

Beyond Cosine Decay: On the effectiveness of Infinite Learning Rate Schedule for Continual Pre-training

The ever-growing availability of unlabeled data presents both opportunities and challenges for training artificial intelligence systems. Whi… (voir plus)le self-supervised learning (SSL) has emerged as a powerful paradigm for extracting meaningful representations from vast amounts of unlabeled data, existing methods still struggle to adapt to the non-stationary, non-IID nature of real-world data streams without forgetting previously learned knowledge. Recent works have adopted a repeated cosine annealing schedule for large-scale continual pre-training; however, these schedules (1) inherently cause forgetting during the re-warming phase and (2) have not been systematically compared to existing continual SSL methods. In this work, we systematically compare the widely used cosine schedule with the recently proposed infinite learning rate schedule and empirically find the latter to be a more effective alternative. Our extensive empirical evaluation across diverse image and language datasets demonstrates that the infinite learning rate schedule consistently enhances continual pre-training performance compared to a repeated cosine decay without being restricted to a fixed iteration budget. For instance, in a small-scale MAE pre-training setup, it outperforms several strong baselines from the literature. We then scale up our experiments to larger MAE pre-training and autoregressive language model pre-training. Our results show that the infinite learning rate schedule remains effective at scale, surpassing repeated cosine decay for both MAE pre-training and zero-shot LM benchmarks.

2025-06-10

ICML.cc/2025/Workshop/ES-FoMo-III (publié)

Model Parallelism With Subnetwork Data Parallelism

Distributed pre-training of large models at scale often imposes heavy memory demands on individual nodes and incurs significant intra-node c… (voir plus)ommunication costs. We propose a novel alternative approach that reduces the memory requirements by training small, structured subnetworks of the model on separate workers. Unlike pipelining, our method avoids inter-node activation communication and maintains bandwidth requirements that are comparable to or lower than standard data parallel communication schemes based on all-reduce. We evaluate two subnetwork construction strategies guided by the principle of ensuring uniform representation of each parameter across the distributed training setup. Our results show that the stochastic block dropping technique consistently outperforms the width-wise subnetwork construction previously explored in federated learning. We empirically attribute this superior performance to stronger gradient alignment in subnetworks that retain blocks having skip connections. Preliminary experiments highlight the promise of our approach, achieving a

2025-06-10

ICML.cc/2025/Workshop/ES-FoMo-III (publié)

MuLoCo: Muon is a practical inner optimizer for DiLoCo

Benjamin Therien

Xiaolong Huang

2025-06-10

ICML.cc/2025/Workshop/ES-FoMo-III (publié)

Geometry-Aware Preference Learning for 3D Texture Generation

Amirhossein Zamani

Tianhao Xie

Amir Aghdam

Tiberiu Popa

Recent advances in 3D generative models have achieved impressive results but 3D contents generated by these models may not align with subjec… (voir plus)tive human preferences or task-specific criteria. Moreover, a core challenge in the 3D texture generation domain remains: most existing approaches rely on repeated calls to 2D text-to-image generative models, which lack an inherent understanding of the 3D structure of the input 3D mesh object. To address this, we propose an end-to-end differentiable preference learning framework that back-propagates human preferences, represented by differentiable reward functions, through the entire 3D generative pipeline, making the process inherently geometry-aware. We demonstrate the effectiveness of our framework using four proposed novel geometry-aware reward functions, offering a more controllable and interpretable pathway for high-quality 3D content creation from natural language.

2025-06-09

ICML.cc/2025/Workshop/MoFA (poster)

Test Time Adaptation Using Adaptive Quantile Recalibration

2025-06-09

ICML.cc/2025/Workshop/PUT (poster)

PyLO: Towards Accessible Learned Optimizers in PyTorch

Paul Janson

Benjamin Therien

Quentin Gregory Anthony

Xiaolong Huang

Abhinav Moudgil

Learned optimizers have been an active research topic over the past decade, with increasing progress toward practical, general-purpose optim… (voir plus)izers that can serve as drop-in replacements for widely used methods like Adam. However, recent advances -- such as VeLO, which was meta-trained for 4000 TPU-months -- remain largely inaccessible to the broader community, in part due to their reliance on JAX and the absence of user-friendly packages for applying the optimizers after meta-training. To address this gap, we introduce PyLO, a PyTorch-based library that brings learned optimizers to the broader machine learning community through familiar, widely adopted workflows. Unlike prior work focused on synthetic or convex tasks, our emphasis is on applying learned optimization to real-world large-scale pre-training tasks. Our release includes a CUDA-accelerated version of the small_fc_lopt learned optimizer architecture from (Metz et al., 2022a), delivering substantial speedups -- from 39.36 to 205.59 samples/sec throughput for training ViT B/16 with batch size 32. PyLO also allows us to easily combine learned optimizers with existing optimization tools such as learning rate schedules and weight decay. When doing so, we find that learned optimizers can substantially benefit. Our code is available at https://github.com/Belilovsky-Lab/pylo

2025-06-08

ICML.cc/2025/Workshop/CODEML (publié)

Unsupervised Test-Time Adaptation for Hepatic Steatosis Grading Using Ultrasound B-Mode Images

Michael Eickenberg

An Tang

Guy Cloutier

Ultrasound (US) is considered a key modality for the clinical assessment of hepatic steatosis (i.e., fatty liver) due to its noninvasiveness… (voir plus) and availability. Deep learning methods have attracted considerable interest in this field, as they are capable of learning patterns in a collection of images and achieve clinically comparable levels of accuracy in steatosis grading. However, variations in patient populations, acquisition protocols, equipment, and operator expertise across clinical sites can introduce domain shifts that reduce model performance when applied outside the original training setting. In response, unsupervised domain adaptation techniques are being investigated to address these shifts, allowing models to generalize more effectively across diverse clinical environments. In this work, we propose a test-time batch normalization (TTN) technique designed to handle domain shift, especially for changes in label distribution, by adapting selected features of batch normalization (BatchNorm) layers in a trained convolutional neural network model. This approach operates in an unsupervised manner, allowing robust adaptation to new distributions without access to label data. The method was evaluated on two abdominal US datasets collected at different institutions, assessing its capability in mitigating domain shift for hepatic steatosis classification. The proposed method reduced the mean absolute error in steatosis grading by 37% and improved the area under the receiver operating characteristic curves (AUC) for steatosis detection from 0.78 to 0.97, compared to nonadapted models. These findings demonstrate the potential of the proposed method to address domain shift in US-based hepatic steatosis diagnosis, minimizing risks associated with deploying trained models in various clinical settings.

2025-03-25

IEEE Transactions on Ultrasonics, Ferroelectrics and Frequency Control (publié)

Adaptive Local Training in Federated Learning.

Donald Shenaj

Pietro Zanuttigh

In federated learning multiple clients collaboratively train a global machine learning model by exchanging their locally trained model weigh… (voir plus)ts instead of raw data. In the standard setting, every client trains its local model for the same number of epochs. We introduce ALT (Adaptive Local Training), a simple yet effective feedback mechanism that can be introduced on top of any federated learning scheme at the client side to limit unnecessary and degrading computations. ALT dynamically adjusts the number of training epochs for each client based on the similarity between the local representation and the global one, ensuring that well-aligned clients can train longer without experiencing client drift while in case of too large drifts the training is stopped earlier. We evaluated ALT on federated partitions of the CIFAR-10 and Tiny-ImageNet datasets, demonstrating its effectiveness in improving both model convergence speed and accuracy. The code is available at https://github.com/LTTM/ALT.

2025-03-04

ICLR.cc/2025/Workshop/MCDC (accepté)