Author name cluster

Jonas Kohler

Possible papers associated with this exact author name in Arrow. This page groups case-insensitive exact name matches and is not a full identity disambiguation profile.

8 papers

2 author rows

AAAI Conference 2025 Conference Paper

Adaptive Guidance: Training-free Acceleration of Conditional Diffusion Models

Angela Castillo
Jonas Kohler
Juan C. Pérez
Juan Pablo Pérez
Albert Pumarola
Bernard Ghanem
Pablo Arbeláez
Ali Thabet

This paper presents a comprehensive study on the role of Classifier-Free Guidance (CFG) in text-conditioned diffusion models from the perspective of inference efficiency. In particular, we relax the default choice of applying CFG in all diffusion steps and instead propose to search for more efficient guidance policies. We formulate the discovery of such policies in the framework of differentiable neural architecture search. Our findings suggest that, as denoising progresses, the updates produced by CFG become increasingly aligned with simple conditional steps, which renders CFG's additional neural network evaluation redundant, especially in the second half of the denoising process. Building upon this insight, we propose "Adaptive Guidance" (AG), an efficient variant of CFG that adaptively omits network evaluations when the denoising process displays convergence. Our experiments demonstrate that AG preserves CFG's image quality while reducing computation by 25%. Thus, AG constitutes a plug-and-play alternative to Guidance Distillation, achieving 50% of the speed-ups of the latter, while being training-free and retaining the capacity to handle negative prompts. We conclude by uncovering further redundancies of CFG in the first half of the diffusion process, showing that entire neural network evaluations can be replaced by simple affine transformations of past score estimates.

PDF Details DOI

ICLR Conference 2025 Conference Paper

Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment

Gregor Bachmann
Sotiris Anagnostidis
Albert Pumarola
Markos Georgopoulos
Artsiom Sanakoyeu
Yuming Du
Edgar Schönfeld
Ali K. Thabet

The performance of large language models (LLMs) is closely linked to their underlying size, leading to ever-growing networks and hence slower inference. Speculative decoding has been proposed as a technique to accelerate autoregressive generation, leveraging a fast draft model to propose candidate tokens, which are then verified in parallel based on their likelihood under the target model. While this approach guarantees to reproduce the target output, it incurs a substantial penalty: many high-quality draft tokens are rejected, even when they represent objectively valid continuations. Indeed, we show that even powerful draft models such as GPT-4o, as well as human text cannot achieve high acceptance rates under the standard verification scheme. This severely limits the speedup potential of current speculative decoding methods, as an early rejection becomes overwhelmingly likely when solely relying on alignment of draft and target. We thus ask the following question: Can we adapt verification to recognize correct, but non-aligned replies? To this end, we draw inspiration from the LLM-as-a-judge framework, which demonstrated that LLMs are able to rate answers in a versatile way. We carefully design a dataset coined TokenCourt to elicit the same capability in the target model by training a compact module on top of the embeddings to produce ``judgements" of the current continuation. We showcase our strategy on the Llama-3.1 family, where our 8B/405B-Judge achieves a speedup of $9\times$ over Llama-405B, while maintaining its quality on a large range of benchmarks. These benefits remain present even in optimized inference frameworks, where our method reaches up to $141$ tokens/s for 8B/70B-Judge and $129$ tokens/s for 8B/405B on $2$ and $8$ H100s respectively.

Details

ICLR Conference 2025 Conference Paper

Storybooth: Training-Free Multi-Subject Consistency for Improved Visual Storytelling

Jaskirat Singh
Junshen K. Chen
Jonas Kohler
Michael F. Cohen

Consistent text-to-image generation depicting the *same* subjects across different images has gained significant recent attention due to its widespread applications in the fields of visual-storytelling and multiple-shot video generation. While remarkable, existing methods often require costly finetuning for each subject and struggle to maintain consistency across multiple characters. In this work, we first analyse the reason for these limitations. Our exploration reveals that the primary-issue stems from *self-attention leakage*, which is exacerbated when trying to ensure consistency across multiple-characters. Motivated by these findings, we next propose a simple yet effective *training and optimization-free approach* for improving multiple-character consistency. In particular, we first leverage multi-modal *chain-of-thought* reasoning in order to *apriori* localize the different subjects across the storyboard frames. The final storyboard images are then generated using a modified diffusion model which includes *1) a bounded cross-attention layer* for ensuring adherence to the initially predicted layout, and *2) a bounded cross-frame self-attention layer* for reducing inter-character attention leakage. Furthermore, we also propose a novel *cross-frame token-merging layer* which allows for improved fine-grain consistency for the storyboard characters. Experimental analysis reveals that proposed approach is not only $\times 30$ faster than prior training-based methods (*eg, textual inversion, dreambooth-lora*) but also surpasses the prior *state-of-the-art*, exhibiting improved multi-character consistency and text-to-image alignment performance.

Details

YNIMG Journal 2024 Journal Article

Deep learning applied to the segmentation of rodent brain MRI data outperforms noisy ground truth on full-fledged brain atlases

Jonas Kohler
Thomas Bielser
Stanislaw Adaszewski
Basil Künnecke
Andreas Bruns

Translational magnetic resonance imaging of the rodent brain provides invaluable information for preclinical drug development. However, the automated segmentation of such images for quantitative analyses is limited compared to human brain imaging mainly due to the inferior anatomical contrast and the resulting less advanced registration and atlasing tools. Here, we investigated the potential of deep learning models for the segmentation of magnetic resonance images of rat brains into an entire set of multiple regions of interest (rather than individual loci), focusing on the development of a robust method that accommodates changes in the input based on differences in animal strain (genotype) and size. Manually generated labels are expensive, so we tested the ability of neural networks to learn brain structures from noisy but inexpensive registration-based labels, allowing very large datasets to be leveraged for training. We compared three distinct model architectures (U-Net, Attention-U-Net and DeepLab) by training them on a dataset of >10, 000 magnetic resonance images of rat brains and found that each model was able to segment the entire brain into predefined sets of 29 and 58 regions, respectively, with the Attention U-Net achieving the best performance. The models canceled out unstructured label noise in the imperfect training data to provide smoother and more symmetric segmentations than registration-based labeling, and were more robust when presented with input variations, thus outperforming the noisy ground truth. Our pipeline also includes uncertainty estimation and an explainability mechanism, hence providing features essential for anomaly detection and quality assurance. In summary, our study shows that deep learning models do achieve accurate brain segmentation in high-throughput quantitative preclinical imaging without the need for expensive expert-generated labels.

Details DOI

NeurIPS Conference 2020 Conference Paper

Batch normalization provably avoids ranks collapse for randomly initialised deep networks

Hadi Daneshmand
Jonas Kohler
Francis Bach
Thomas Hofmann
Aurelien Lucchi

Randomly initialized neural networks are known to become harder to train with increasing depth, unless architectural enhancements like residual connections and batch normalization are used. We here investigate this phenomenon by revisiting the connection between random initialization in deep networks and spectral instabilities in products of random matrices. Given the rich literature on random matrices, it is not surprising to find that the rank of the intermediate representations in unnormalized networks collapses quickly with depth. In this work we highlight the fact that batch normalization is an effective strategy to avoid rank collapse for both linear and ReLU networks. Leveraging tools from Markov chain theory, we derive a meaningful lower rank bound in deep linear networks. Empirically, we also demonstrate that this rank robustness generalizes to ReLU nets. Finally, we conduct an extensive set of experiments on real-world data sets, which confirm that rank stability is indeed a crucial condition for training modern-day deep neural architectures.

PDF Details

UAI Conference 2019 Conference Paper

The Role of Memory in Stochastic Optimization

Antonio Orvieto
Jonas Kohler
Aurélien Lucchi

The choice of how to retain information about past gradients dramatically affects the convergence properties of state-of-the-art stochastic optimization methods, such as Heavy-ball, Nesterov’s momentum, RMSprop and Adam. Building on this observation, we use stochastic differential equations (SDEs) to explicitly study the role of memory in gradient-based algorithms. We first derive a general continuous-time model that can incorporate arbitrary types of memory, for both deterministic and stochastic settings. We provide convergence guarantees for this SDE for weakly-quasi-convex and quadratically growing functions. We then demonstrate how to discretize this SDE to get a flexible discrete-time algorithm that can implement a board spectrum of memories ranging from short- to long-term. Not only does this algorithm increase the degrees of freedom in algorithmic choice for practitioners but it also comes with better stability properties than classical momentum in the convex stochastic setting. In particular, no iterate averaging is needed for convergence. Interestingly, our analysis also provides a novel interpretation of Nesterov’s momentum as stable gradient amplification and highlights a possible reason for its unstable behavior in the (convex) stochastic setting. Furthermore, we discuss the use of long term memory for second-moment estimation in adaptive methods, such as Adam and RMSprop. Finally, we provide an extensive experimental study of the effect of different types of memory in both convex and nonconvex settings.

Details

ICML Conference 2018 Conference Paper

Escaping Saddles with Stochastic Gradients

Hadi Daneshmand
Jonas Kohler
Aurélien Lucchi
Thomas Hofmann 0001

We analyze the variance of stochastic gradients along negative curvature directions in certain non-convex machine learning models and show that stochastic gradients indeed exhibit a strong component along these directions. Furthermore, we show that - contrary to the case of isotropic noise - this variance is proportional to the magnitude of the corresponding eigenvalues and not decreasing in the dimensionality. Based upon this bservation we propose a new assumption under which we show that the injection of explicit, isotropic noise usually applied to make gradient descent escape saddle points can successfully be replaced by a simple SGD step. Additionally - and under the same condition - we derive the first convergence rate for plain SGD to a second-order stationary point in a number of iterations that is independent of the problem dimension.

Details

ICML Conference 2017 Conference Paper

Sub-sampled Cubic Regularization for Non-convex Optimization

Jonas Kohler
Aurélien Lucchi

We consider the minimization of non-convex functions that typically arise in machine learning. Specifically, we focus our attention on a variant of trust region methods known as cubic regularization. This approach is particularly attractive because it escapes strict saddle points and it provides stronger convergence guarantees than first- and second-order as well as classical trust region methods. However, it suffers from a high computational complexity that makes it impractical for large-scale learning. Here, we propose a novel method that uses sub-sampling to lower this computational cost. By the use of concentration inequalities we provide a sampling scheme that gives sufficiently accurate gradient and Hessian approximations to retain the strong global and local convergence guarantees of cubically regularized methods. To the best of our knowledge this is the first work that gives global convergence guarantees for a sub-sampled variant of cubic regularization on non-convex functions. Furthermore, we provide experimental results supporting our theory.

Details