-
Conditioning diffusion models by explicit forward-backward bridging
Authors:
Adrien Corenflos,
Zheng Zhao,
Simo Särkkä,
Jens Sjölund,
Thomas B. Schön
Abstract:
Given an unconditional diffusion model $π(x, y)$, using it to perform conditional simulation $π(x \mid y)$ is still largely an open question and is typically achieved by learning conditional drifts to the denoising SDE after the fact. In this work, we express conditional simulation as an inference problem on an augmented space corresponding to a partial SDE bridge. This perspective allows us to im…
▽ More
Given an unconditional diffusion model $π(x, y)$, using it to perform conditional simulation $π(x \mid y)$ is still largely an open question and is typically achieved by learning conditional drifts to the denoising SDE after the fact. In this work, we express conditional simulation as an inference problem on an augmented space corresponding to a partial SDE bridge. This perspective allows us to implement efficient and principled particle Gibbs and pseudo-marginal samplers marginally targeting the conditional distribution $π(x \mid y)$. Contrary to existing methodology, our methods do not introduce any additional approximation to the unconditional diffusion model aside from the Monte Carlo error. We showcase the benefits and drawbacks of our approach on a series of synthetic and real data examples.
△ Less
Submitted 22 May, 2024;
originally announced May 2024.
-
Efficient Radiation Treatment Planning based on Voxel Importance
Authors:
Sebastian Mair,
Anqi Fu,
Jens Sjölund
Abstract:
Optimization is a time-consuming part of radiation treatment planning. We propose to reduce the optimization problem by only using a representative subset of informative voxels. This way, we improve planning efficiency while maintaining or enhancing the plan quality. To reduce the computational complexity of the optimization problem, we propose to subsample the set of voxels via importance samplin…
▽ More
Optimization is a time-consuming part of radiation treatment planning. We propose to reduce the optimization problem by only using a representative subset of informative voxels. This way, we improve planning efficiency while maintaining or enhancing the plan quality. To reduce the computational complexity of the optimization problem, we propose to subsample the set of voxels via importance sampling. We derive a sampling distribution based on an importance score that we obtain from pre-solving an easy optimization problem involving a simplified probing objective. By solving a reduced version of the original optimization problem using this subset, we effectively reduce the problem's size and computational demands while accounting for regions in which satisfactory dose deliveries are challenging. In contrast to other stochastic (sub-)sampling methods, our technique only requires a single sampling step to define a reduced optimization problem. This problem can be efficiently solved using established solvers. Empirical experiments on open benchmark data highlight substantially reduced optimization times, up to 50 times faster than the original ones, for intensity-modulated radiation therapy (IMRT), all while upholding plan quality comparable to traditional methods. Our approach has the potential to significantly accelerate radiation treatment planning by addressing its inherent computational challenges. We reduce the treatment planning time by reducing the size of the optimization problem rather than improving the optimization method. Our efforts are thus complementary to much of the previous developments.
△ Less
Submitted 6 May, 2024;
originally announced May 2024.
-
Photo-Realistic Image Restoration in the Wild with Controlled Vision-Language Models
Authors:
Ziwei Luo,
Fredrik K. Gustafsson,
Zheng Zhao,
Jens Sjölund,
Thomas B. Schön
Abstract:
Though diffusion models have been successfully applied to various image restoration (IR) tasks, their performance is sensitive to the choice of training datasets. Typically, diffusion models trained in specific datasets fail to recover images that have out-of-distribution degradations. To address this problem, this work leverages a capable vision-language model and a synthetic degradation pipeline…
▽ More
Though diffusion models have been successfully applied to various image restoration (IR) tasks, their performance is sensitive to the choice of training datasets. Typically, diffusion models trained in specific datasets fail to recover images that have out-of-distribution degradations. To address this problem, this work leverages a capable vision-language model and a synthetic degradation pipeline to learn image restoration in the wild (wild IR). More specifically, all low-quality images are simulated with a synthetic degradation pipeline that contains multiple common degradations such as blur, resize, noise, and JPEG compression. Then we introduce robust training for a degradation-aware CLIP model to extract enriched image content features to assist high-quality image restoration. Our base diffusion model is the image restoration SDE (IR-SDE). Built upon it, we further present a posterior sampling strategy for fast noise-free image generation. We evaluate our model on both synthetic and real-world degradation datasets. Moreover, experiments on the unified image restoration task illustrate that the proposed posterior sampling improves image generation quality for various degradations.
△ Less
Submitted 15 April, 2024;
originally announced April 2024.
-
Ising on the Graph: Task-specific Graph Subsampling via the Ising Model
Authors:
Maria Bånkestad,
Jennifer Andersson,
Sebastian Mair,
Jens Sjölund
Abstract:
Reducing a graph while preserving its overall structure is an important problem with many applications. Typically, the reduction approaches either remove edges (sparsification) or merge nodes (coarsening) in an unsupervised way with no specific downstream task in mind. In this paper, we present an approach for subsampling graph structures using an Ising model defined on either the nodes or edges a…
▽ More
Reducing a graph while preserving its overall structure is an important problem with many applications. Typically, the reduction approaches either remove edges (sparsification) or merge nodes (coarsening) in an unsupervised way with no specific downstream task in mind. In this paper, we present an approach for subsampling graph structures using an Ising model defined on either the nodes or edges and learning the external magnetic field of the Ising model using a graph neural network. Our approach is task-specific as it can learn how to reduce a graph for a specific downstream task in an end-to-end fashion. The utilized loss function of the task does not even have to be differentiable. We showcase the versatility of our approach on three distinct applications: image segmentation, 3D shape sparsification, and sparse approximate matrix inverse determination.
△ Less
Submitted 15 February, 2024;
originally announced February 2024.
-
Entropy-regularized Diffusion Policy with Q-Ensembles for Offline Reinforcement Learning
Authors:
Ruoqi Zhang,
Ziwei Luo,
Jens Sjölund,
Thomas B. Schön,
Per Mattsson
Abstract:
This paper presents advanced techniques of training diffusion policies for offline reinforcement learning (RL). At the core is a mean-reverting stochastic differential equation (SDE) that transfers a complex action distribution into a standard Gaussian and then samples actions conditioned on the environment state with a corresponding reverse-time SDE, like a typical diffusion policy. We show that…
▽ More
This paper presents advanced techniques of training diffusion policies for offline reinforcement learning (RL). At the core is a mean-reverting stochastic differential equation (SDE) that transfers a complex action distribution into a standard Gaussian and then samples actions conditioned on the environment state with a corresponding reverse-time SDE, like a typical diffusion policy. We show that such an SDE has a solution that we can use to calculate the log probability of the policy, yielding an entropy regularizer that improves the exploration of offline datasets. To mitigate the impact of inaccurate value functions from out-of-distribution data points, we further propose to learn the lower confidence bound of Q-ensembles for more robust policy improvement. By combining the entropy-regularized diffusion policy with Q-ensembles in offline RL, our method achieves state-of-the-art performance on most tasks in D4RL benchmarks. Code is available at \href{https://github.com/ruoqizzz/Entropy-Regularized-Diffusion-Policy-with-QEnsemble}{https://github.com/ruoqizzz/Entropy-Regularized-Diffusion-Policy-with-QEnsemble}.
△ Less
Submitted 6 February, 2024;
originally announced February 2024.
-
Variational Elliptical Processes
Authors:
Maria Bånkestad,
Jens Sjölund,
Jalil Taghia,
Thomas B. Schöon
Abstract:
We present elliptical processes, a family of non-parametric probabilistic models that subsume Gaussian processes and Student's t processes. This generalization includes a range of new heavy-tailed behaviors while retaining computational tractability. Elliptical processes are based on a representation of elliptical distributions as a continuous mixture of Gaussian distributions. We parameterize thi…
▽ More
We present elliptical processes, a family of non-parametric probabilistic models that subsume Gaussian processes and Student's t processes. This generalization includes a range of new heavy-tailed behaviors while retaining computational tractability. Elliptical processes are based on a representation of elliptical distributions as a continuous mixture of Gaussian distributions. We parameterize this mixture distribution as a spline normalizing flow, which we train using variational inference. The proposed form of the variational posterior enables a sparse variational elliptical process applicable to large-scale problems. We highlight advantages compared to Gaussian processes through regression and classification experiments. Elliptical processes can supersede Gaussian processes in several settings, including cases where the likelihood is non-Gaussian or when accurate tail modeling is essential.
△ Less
Submitted 21 November, 2023;
originally announced November 2023.
-
On Feynman--Kac training of partial Bayesian neural networks
Authors:
Zheng Zhao,
Sebastian Mair,
Thomas B. Schön,
Jens Sjölund
Abstract:
Recently, partial Bayesian neural networks (pBNNs), which only consider a subset of the parameters to be stochastic, were shown to perform competitively with full Bayesian neural networks. However, pBNNs are often multi-modal in the latent variable space and thus challenging to approximate with parametric models. To address this problem, we propose an efficient sampling-based training strategy, wh…
▽ More
Recently, partial Bayesian neural networks (pBNNs), which only consider a subset of the parameters to be stochastic, were shown to perform competitively with full Bayesian neural networks. However, pBNNs are often multi-modal in the latent variable space and thus challenging to approximate with parametric models. To address this problem, we propose an efficient sampling-based training strategy, wherein the training of a pBNN is formulated as simulating a Feynman--Kac model. We then describe variations of sequential Monte Carlo samplers that allow us to simultaneously estimate the parameters and the latent posterior distribution of this model at a tractable computational cost. Using various synthetic and real-world datasets we show that our proposed training scheme outperforms the state of the art in terms of predictive performance.
△ Less
Submitted 27 February, 2024; v1 submitted 30 October, 2023;
originally announced October 2023.
-
Controlling Vision-Language Models for Multi-Task Image Restoration
Authors:
Ziwei Luo,
Fredrik K. Gustafsson,
Zheng Zhao,
Jens Sjölund,
Thomas B. Schön
Abstract:
Vision-language models such as CLIP have shown great impact on diverse downstream tasks for zero-shot or label-free predictions. However, when it comes to low-level vision such as image restoration their performance deteriorates dramatically due to corrupted inputs. In this paper, we present a degradation-aware vision-language model (DA-CLIP) to better transfer pretrained vision-language models to…
▽ More
Vision-language models such as CLIP have shown great impact on diverse downstream tasks for zero-shot or label-free predictions. However, when it comes to low-level vision such as image restoration their performance deteriorates dramatically due to corrupted inputs. In this paper, we present a degradation-aware vision-language model (DA-CLIP) to better transfer pretrained vision-language models to low-level vision tasks as a multi-task framework for image restoration. More specifically, DA-CLIP trains an additional controller that adapts the fixed CLIP image encoder to predict high-quality feature embeddings. By integrating the embedding into an image restoration network via cross-attention, we are able to pilot the model to learn a high-fidelity image reconstruction. The controller itself will also output a degradation feature that matches the real corruptions of the input, yielding a natural classifier for different degradation types. In addition, we construct a mixed degradation dataset with synthetic captions for DA-CLIP training. Our approach advances state-of-the-art performance on both \emph{degradation-specific} and \emph{unified} image restoration tasks, showing a promising direction of prompting image restoration with large-scale pretrained vision-language models. Our code is available at https://github.com/Algolzw/daclip-uir.
△ Less
Submitted 28 February, 2024; v1 submitted 2 October, 2023;
originally announced October 2023.
-
ICML 2023 Topological Deep Learning Challenge : Design and Results
Authors:
Mathilde Papillon,
Mustafa Hajij,
Helen Jenne,
Johan Mathe,
Audun Myers,
Theodore Papamarkou,
Tolga Birdal,
Tamal Dey,
Tim Doster,
Tegan Emerson,
Gurusankar Gopalakrishnan,
Devendra Govil,
Aldo Guzmán-Sáenz,
Henry Kvinge,
Neal Livesay,
Soham Mukherjee,
Shreyas N. Samaga,
Karthikeyan Natesan Ramamurthy,
Maneel Reddy Karri,
Paul Rosen,
Sophia Sanborn,
Robin Walters,
Jens Agerberg,
Sadrodin Barikbin,
Claudio Battiloro
, et al. (31 additional authors not shown)
Abstract:
This paper presents the computational challenge on topological deep learning that was hosted within the ICML 2023 Workshop on Topology and Geometry in Machine Learning. The competition asked participants to provide open-source implementations of topological neural networks from the literature by contributing to the python packages TopoNetX (data processing) and TopoModelX (deep learning). The chal…
▽ More
This paper presents the computational challenge on topological deep learning that was hosted within the ICML 2023 Workshop on Topology and Geometry in Machine Learning. The competition asked participants to provide open-source implementations of topological neural networks from the literature by contributing to the python packages TopoNetX (data processing) and TopoModelX (deep learning). The challenge attracted twenty-eight qualifying submissions in its two-month duration. This paper describes the design of the challenge and summarizes its main findings.
△ Less
Submitted 18 January, 2024; v1 submitted 26 September, 2023;
originally announced September 2023.
-
Personalized Privacy Amplification via Importance Sampling
Authors:
Dominik Fay,
Sebastian Mair,
Jens Sjölund
Abstract:
We examine the privacy-enhancing properties of importance sampling. In importance sampling, selection probabilities are heterogeneous and each selected data point is weighted by the reciprocal of its selection probability. Due to the heterogeneity of importance sampling, we express our results within the framework of personalized differential privacy. We first consider the general case where an ar…
▽ More
We examine the privacy-enhancing properties of importance sampling. In importance sampling, selection probabilities are heterogeneous and each selected data point is weighted by the reciprocal of its selection probability. Due to the heterogeneity of importance sampling, we express our results within the framework of personalized differential privacy. We first consider the general case where an arbitrary personalized differentially private mechanism is subsampled with an arbitrary importance sampling distribution and show that the resulting mechanism also satisfies personalized differential privacy. This constitutes an extension of the established privacy amplification by subsampling result to importance sampling. Then, for any fixed mechanism, we derive the sampling distribution that achieves the optimal sampling rate subject to a worst-case privacy constraint. Empirically, we evaluate the privacy, efficiency, and accuracy of importance sampling on the example of k-means clustering.
△ Less
Submitted 15 February, 2024; v1 submitted 5 July, 2023;
originally announced July 2023.
-
Risk-sensitive Actor-free Policy via Convex Optimization
Authors:
Ruoqi Zhang,
Jens Sjölund
Abstract:
Traditional reinforcement learning methods optimize agents without considering safety, potentially resulting in unintended consequences. In this paper, we propose an optimal actor-free policy that optimizes a risk-sensitive criterion based on the conditional value at risk. The risk-sensitive objective function is modeled using an input-convex neural network ensuring convexity with respect to the a…
▽ More
Traditional reinforcement learning methods optimize agents without considering safety, potentially resulting in unintended consequences. In this paper, we propose an optimal actor-free policy that optimizes a risk-sensitive criterion based on the conditional value at risk. The risk-sensitive objective function is modeled using an input-convex neural network ensuring convexity with respect to the actions and enabling the identification of globally optimal actions through simple gradient-following methods. Experimental results demonstrate the efficacy of our approach in maintaining effective risk control.
△ Less
Submitted 30 June, 2023;
originally announced July 2023.
-
Neural incomplete factorization: learning preconditioners for the conjugate gradient method
Authors:
Paul Häusner,
Ozan Öktem,
Jens Sjölund
Abstract:
Finding suitable preconditioners to accelerate iterative solution methods, such as the conjugate gradient method, is an active area of research. In this paper, we develop a computationally efficient data-driven approach to replace the typically hand-engineered algorithms with neural networks. Optimizing the condition number of the linear system directly is computationally infeasible. Instead, our…
▽ More
Finding suitable preconditioners to accelerate iterative solution methods, such as the conjugate gradient method, is an active area of research. In this paper, we develop a computationally efficient data-driven approach to replace the typically hand-engineered algorithms with neural networks. Optimizing the condition number of the linear system directly is computationally infeasible. Instead, our method generates an incomplete factorization of the matrix and is, therefore, referred to as neural incomplete factorization (NeuralIF). For efficient training, we utilize a stochastic approximation of the Frobenius loss which only requires matrix-vector multiplications. At the core of our method is a novel messagepassing block, inspired by sparse matrix theory, that aligns with the objective of finding a sparse factorization of the matrix. By replacing conventional preconditioners used within the conjugate gradient method by data-driven models based on graph neural networks, we accelerate the iterative solving procedure. We evaluate our proposed method on both a synthetic and a real-world problem arising from scientific computing and show its ability to reduce the solving time while remaining computationally efficient.
△ Less
Submitted 5 February, 2024; v1 submitted 25 May, 2023;
originally announced May 2023.
-
Refusion: Enabling Large-Size Realistic Image Restoration with Latent-Space Diffusion Models
Authors:
Ziwei Luo,
Fredrik K. Gustafsson,
Zheng Zhao,
Jens Sjölund,
Thomas B. Schön
Abstract:
This work aims to improve the applicability of diffusion models in realistic image restoration. Specifically, we enhance the diffusion model in several aspects such as network architecture, noise level, denoising steps, training image size, and optimizer/scheduler. We show that tuning these hyperparameters allows us to achieve better performance on both distortion and perceptual scores. We also pr…
▽ More
This work aims to improve the applicability of diffusion models in realistic image restoration. Specifically, we enhance the diffusion model in several aspects such as network architecture, noise level, denoising steps, training image size, and optimizer/scheduler. We show that tuning these hyperparameters allows us to achieve better performance on both distortion and perceptual scores. We also propose a U-Net based latent diffusion model which performs diffusion in a low-resolution latent space while preserving high-resolution information from the original input for the decoding process. Compared to the previous latent-diffusion model which trains a VAE-GAN to compress the image, our proposed U-Net compression strategy is significantly more stable and can recover highly accurate images without relying on adversarial optimization. Importantly, these modifications allow us to apply diffusion models to various image restoration tasks, including real-world shadow removal, HR non-homogeneous dehazing, stereo super-resolution, and bokeh effect transformation. By simply replacing the datasets and slightly changing the noise network, our model, named Refusion, is able to deal with large-size images (e.g., 6000 x 4000 x 3 in HR dehazing) and produces good results on all the above restoration problems. Our Refusion achieves the best perceptual performance in the NTIRE 2023 Image Shadow Removal Challenge and wins 2nd place overall.
△ Less
Submitted 17 April, 2023;
originally announced April 2023.
-
Archetypal Analysis++: Rethinking the Initialization Strategy
Authors:
Sebastian Mair,
Jens Sjölund
Abstract:
Archetypal analysis is a matrix factorization method with convexity constraints. Due to local minima, a good initialization is essential, but frequently used initialization methods yield either sub-optimal starting points or are prone to get stuck in poor local minima. In this paper, we propose archetypal analysis++ (AA++), a probabilistic initialization strategy for archetypal analysis that seque…
▽ More
Archetypal analysis is a matrix factorization method with convexity constraints. Due to local minima, a good initialization is essential, but frequently used initialization methods yield either sub-optimal starting points or are prone to get stuck in poor local minima. In this paper, we propose archetypal analysis++ (AA++), a probabilistic initialization strategy for archetypal analysis that sequentially samples points based on their influence on the objective function, similar to $k$-means++. In fact, we argue that $k$-means++ already approximates the proposed initialization method. Furthermore, we suggest to adapt an efficient Monte Carlo approximation of $k$-means++ to AA++. In an extensive empirical evaluation of 15 real-world data sets of varying sizes and dimensionalities and considering two pre-processing strategies, we show that AA++ almost always outperforms all baselines, including the most frequently used ones.
△ Less
Submitted 13 May, 2024; v1 submitted 31 January, 2023;
originally announced January 2023.
-
Image Restoration with Mean-Reverting Stochastic Differential Equations
Authors:
Ziwei Luo,
Fredrik K. Gustafsson,
Zheng Zhao,
Jens Sjölund,
Thomas B. Schön
Abstract:
This paper presents a stochastic differential equation (SDE) approach for general-purpose image restoration. The key construction consists in a mean-reverting SDE that transforms a high-quality image into a degraded counterpart as a mean state with fixed Gaussian noise. Then, by simulating the corresponding reverse-time SDE, we are able to restore the origin of the low-quality image without relyin…
▽ More
This paper presents a stochastic differential equation (SDE) approach for general-purpose image restoration. The key construction consists in a mean-reverting SDE that transforms a high-quality image into a degraded counterpart as a mean state with fixed Gaussian noise. Then, by simulating the corresponding reverse-time SDE, we are able to restore the origin of the low-quality image without relying on any task-specific prior knowledge. Crucially, the proposed mean-reverting SDE has a closed-form solution, allowing us to compute the ground truth time-dependent score and learn it with a neural network. Moreover, we propose a maximum likelihood objective to learn an optimal reverse trajectory that stabilizes the training and improves the restoration results. The experiments show that our proposed method achieves highly competitive performance in quantitative comparisons on image deraining, deblurring, and denoising, setting a new state-of-the-art on two deraining datasets. Finally, the general applicability of our approach is further demonstrated via qualitative results on image super-resolution, inpainting, and dehazing. Code is available at https://github.com/Algolzw/image-restoration-sde.
△ Less
Submitted 31 May, 2023; v1 submitted 27 January, 2023;
originally announced January 2023.
-
A Tutorial on Parametric Variational Inference
Authors:
Jens Sjölund
Abstract:
Variational inference uses optimization, rather than integration, to approximate the marginal likelihood, and thereby the posterior, in a Bayesian model. Thanks to advances in computational scalability made in the last decade, variational inference is now the preferred choice for many high-dimensional models and large datasets. This tutorial introduces variational inference from the parametric per…
▽ More
Variational inference uses optimization, rather than integration, to approximate the marginal likelihood, and thereby the posterior, in a Bayesian model. Thanks to advances in computational scalability made in the last decade, variational inference is now the preferred choice for many high-dimensional models and large datasets. This tutorial introduces variational inference from the parametric perspective that dominates these recent developments, in contrast to the mean-field perspective commonly found in other introductory texts.
△ Less
Submitted 3 January, 2023;
originally announced January 2023.
-
Probabilistic Estimation of Instantaneous Frequencies of Chirp Signals
Authors:
Zheng Zhao,
Simo Särkkä,
Jens Sjölund,
Thomas B. Schön
Abstract:
We present a continuous-time probabilistic approach for estimating the chirp signal and its instantaneous frequency function when the true forms of these functions are not accessible. Our model represents these functions by non-linearly cascaded Gaussian processes represented as non-linear stochastic differential equations. The posterior distribution of the functions is then estimated with stochas…
▽ More
We present a continuous-time probabilistic approach for estimating the chirp signal and its instantaneous frequency function when the true forms of these functions are not accessible. Our model represents these functions by non-linearly cascaded Gaussian processes represented as non-linear stochastic differential equations. The posterior distribution of the functions is then estimated with stochastic filters and smoothers. We compute a (posterior) Cramér--Rao lower bound for the Gaussian process model, and derive a theoretical upper bound for the estimation error in the mean squared sense. The experiments show that the proposed method outperforms a number of state-of-the-art methods on a synthetic data. We also show that the method works out-of-the-box for two real-world datasets.
△ Less
Submitted 13 February, 2023; v1 submitted 12 May, 2022;
originally announced May 2022.
-
NUQ: A Noise Metric for Diffusion MRI via Uncertainty Discrepancy Quantification
Authors:
Shreyas Fadnavis,
Jens Sjölund,
Anders Eklund,
Eleftherios Garyfallidis
Abstract:
Diffusion MRI (dMRI) is the only non-invasive technique sensitive to tissue micro-architecture, which can, in turn, be used to reconstruct tissue microstructure and white matter pathways. The accuracy of such tasks is hampered by the low signal-to-noise ratio in dMRI. Today, the noise is characterized mainly by visual inspection of residual maps and estimated standard deviation. However, it is har…
▽ More
Diffusion MRI (dMRI) is the only non-invasive technique sensitive to tissue micro-architecture, which can, in turn, be used to reconstruct tissue microstructure and white matter pathways. The accuracy of such tasks is hampered by the low signal-to-noise ratio in dMRI. Today, the noise is characterized mainly by visual inspection of residual maps and estimated standard deviation. However, it is hard to estimate the impact of noise on downstream tasks based only on such qualitative assessments. To address this issue, we introduce a novel metric, Noise Uncertainty Quantification (NUQ), for quantitative image quality analysis in the absence of a ground truth reference image. NUQ uses a recent Bayesian formulation of dMRI models to estimate the uncertainty of microstructural measures. Specifically, NUQ uses the maximum mean discrepancy metric to compute a pooled quality score by comparing samples drawn from the posterior distribution of the microstructure measures. We show that NUQ allows a fine-grained analysis of noise, capturing details that are visually imperceptible. We perform qualitative and quantitative comparisons on real datasets, showing that NUQ generates consistent scores across different denoisers and acquisitions. Lastly, by using NUQ on a cohort of schizophrenics and controls, we quantify the substantial impact of denoising on group differences.
△ Less
Submitted 4 March, 2022; v1 submitted 3 March, 2022;
originally announced March 2022.
-
Graph-based Neural Acceleration for Nonnegative Matrix Factorization
Authors:
Jens Sjölund,
Maria Bånkestad
Abstract:
We describe a graph-based neural acceleration technique for nonnegative matrix factorization that builds upon a connection between matrices and bipartite graphs that is well-known in certain fields, e.g., sparse linear algebra, but has not yet been exploited to design graph neural networks for matrix computations. We first consider low-rank factorization more broadly and propose a graph representa…
▽ More
We describe a graph-based neural acceleration technique for nonnegative matrix factorization that builds upon a connection between matrices and bipartite graphs that is well-known in certain fields, e.g., sparse linear algebra, but has not yet been exploited to design graph neural networks for matrix computations. We first consider low-rank factorization more broadly and propose a graph representation of the problem suited for graph neural networks. Then, we focus on the task of nonnegative matrix factorization and propose a graph neural network that interleaves bipartite self-attention layers with updates based on the alternating direction method of multipliers. Our empirical evaluation on synthetic and two real-world datasets shows that we attain substantial acceleration, even though we only train in an unsupervised fashion on smaller synthetic instances.
△ Less
Submitted 1 February, 2022;
originally announced February 2022.
-
Learn2Reg: comprehensive multi-task medical image registration challenge, dataset and evaluation in the era of deep learning
Authors:
Alessa Hering,
Lasse Hansen,
Tony C. W. Mok,
Albert C. S. Chung,
Hanna Siebert,
Stephanie Häger,
Annkristin Lange,
Sven Kuckertz,
Stefan Heldmann,
Wei Shao,
Sulaiman Vesal,
Mirabela Rusu,
Geoffrey Sonn,
Théo Estienne,
Maria Vakalopoulou,
Luyi Han,
Yunzhi Huang,
Pew-Thian Yap,
Mikael Brudfors,
Yaël Balbastre,
Samuel Joutard,
Marc Modat,
Gal Lifshitz,
Dan Raviv,
**xin Lv
, et al. (28 additional authors not shown)
Abstract:
Image registration is a fundamental medical image analysis task, and a wide variety of approaches have been proposed. However, only a few studies have comprehensively compared medical image registration approaches on a wide range of clinically relevant tasks. This limits the development of registration methods, the adoption of research advances into practice, and a fair benchmark across competing…
▽ More
Image registration is a fundamental medical image analysis task, and a wide variety of approaches have been proposed. However, only a few studies have comprehensively compared medical image registration approaches on a wide range of clinically relevant tasks. This limits the development of registration methods, the adoption of research advances into practice, and a fair benchmark across competing approaches. The Learn2Reg challenge addresses these limitations by providing a multi-task medical image registration data set for comprehensive characterisation of deformable registration algorithms. A continuous evaluation will be possible at https://learn2reg.grand-challenge.org. Learn2Reg covers a wide range of anatomies (brain, abdomen, and thorax), modalities (ultrasound, CT, MR), availability of annotations, as well as intra- and inter-patient registration evaluation. We established an easily accessible framework for training and validation of 3D registration methods, which enabled the compilation of results of over 65 individual method submissions from more than 20 unique teams. We used a complementary set of metrics, including robustness, accuracy, plausibility, and runtime, enabling unique insight into the current state-of-the-art of medical image registration. This paper describes datasets, tasks, evaluation methods and results of the challenge, as well as results of further analysis of transferability to new datasets, the importance of label supervision, and resulting bias. While no single approach worked best across all tasks, many methodological aspects could be identified that push the performance of medical image registration to new state-of-the-art performance. Furthermore, we demystified the common belief that conventional registration methods have to be much slower than deep-learning-based methods.
△ Less
Submitted 7 October, 2022; v1 submitted 8 December, 2021;
originally announced December 2021.
-
Unsupervised dynamic modeling of medical image transformation
Authors:
Niklas Gunnarsson,
Peter Kimstrand,
Jens Sjölund,
Thomas B. Schön
Abstract:
Spatiotemporal imaging has applications in e.g. cardiac diagnostics, surgical guidance, and radiotherapy monitoring, In this paper, we explain the temporal motion by identifying the underlying dynamics, only based on the sequential images. Our dynamical model maps the inputs of observed high-dimensional sequential images to a low-dimensional latent space wherein a linear relationship between a hid…
▽ More
Spatiotemporal imaging has applications in e.g. cardiac diagnostics, surgical guidance, and radiotherapy monitoring, In this paper, we explain the temporal motion by identifying the underlying dynamics, only based on the sequential images. Our dynamical model maps the inputs of observed high-dimensional sequential images to a low-dimensional latent space wherein a linear relationship between a hidden state process and the lower-dimensional representation of the inputs holds. For this, we use a conditional variational auto-encoder (CVAE) to nonlinearly map the higher-dimensional image to a lower-dimensional space, wherein we model the dynamics with a linear Gaussian state-space model (LG-SSM). The model, a modified version of the Kalman variational auto-encoder, is end-to-end trainable, and the weights, both in the CVAE and LG-SSM, are simultaneously updated by maximizing the evidence lower bound of the marginal likelihood. In contrast to the original model, we explain the motion with a spatial transformation from one image to another. This results in sharper reconstructions and the possibility of transferring auxiliary information, such as segmentation, through the image sequence. Our experiments, on cardiac ultrasound time series, show that the dynamic model outperforms traditional image registration in execution time, to a similar performance. Further, our model offers the possibility to impute and extrapolate for missing samples.
△ Less
Submitted 7 November, 2022; v1 submitted 1 March, 2021;
originally announced March 2021.
-
Decentralized Differentially Private Segmentation with PATE
Authors:
Dominik Fay,
Jens Sjölund,
Tobias J. Oechtering
Abstract:
When it comes to preserving privacy in medical machine learning, two important considerations are (1) kee** data local to the institution and (2) avoiding inference of sensitive information from the trained model. These are often addressed using federated learning and differential privacy, respectively. However, the commonly used Federated Averaging algorithm requires a high degree of synchroniz…
▽ More
When it comes to preserving privacy in medical machine learning, two important considerations are (1) kee** data local to the institution and (2) avoiding inference of sensitive information from the trained model. These are often addressed using federated learning and differential privacy, respectively. However, the commonly used Federated Averaging algorithm requires a high degree of synchronization between participating institutions. For this reason, we turn our attention to Private Aggregation of Teacher Ensembles (PATE), where all local models can be trained independently without inter-institutional communication. The purpose of this paper is thus to explore how PATE -- originally designed for classification -- can best be adapted for semantic segmentation. To this end, we build low-dimensional representations of segmentation masks which the student can obtain through low-sensitivity queries to the private aggregator. On the Brain Tumor Segmentation (BraTS 2019) dataset, an Autoencoder-based PATE variant achieves a higher Dice coefficient for the same privacy guarantee than prior work based on noisy Federated Averaging.
△ Less
Submitted 9 April, 2020;
originally announced April 2020.
-
Registration by tracking for sequential 2D MRI
Authors:
Niklas Gunnarsson,
Jens Sjölund,
Thomas B. Schön
Abstract:
Our anatomy is in constant motion. With modern MR imaging it is possible to record this motion in real-time during an ongoing radiation therapy session. In this paper we present an image registration method that exploits the sequential nature of 2D MR images to estimate the corresponding displacement field. The method employs several discriminative correlation filters that independently track spec…
▽ More
Our anatomy is in constant motion. With modern MR imaging it is possible to record this motion in real-time during an ongoing radiation therapy session. In this paper we present an image registration method that exploits the sequential nature of 2D MR images to estimate the corresponding displacement field. The method employs several discriminative correlation filters that independently track specific points. Together with a sparse-to-dense interpolation scheme we can then estimate of the displacement field. The discriminative correlation filters are trained online, and our method is modality agnostic. For the interpolation scheme we use a neural network with normalized convolutions that is trained using synthetic diffeomorphic displacement fields. The method is evaluated on a segmented cardiac dataset and when compared to two conventional methods we observe an improved performance. This improvement is especially pronounced when it comes to the detection of larger motions of small objects.
△ Less
Submitted 24 March, 2020;
originally announced March 2020.
-
The Elliptical Processes: a Family of Fat-tailed Stochastic Processes
Authors:
Maria Bånkestad,
Jens Sjölund,
Jalil Taghia,
Thomas Schön
Abstract:
We present the elliptical processes -- a family of non-parametric probabilistic models that subsumes the Gaussian process and the Student-t process. This generalization includes a range of new fat-tailed behaviors yet retains computational tractability. We base the elliptical processes on a representation of elliptical distributions as a continuous mixture of Gaussian distributions and derive clos…
▽ More
We present the elliptical processes -- a family of non-parametric probabilistic models that subsumes the Gaussian process and the Student-t process. This generalization includes a range of new fat-tailed behaviors yet retains computational tractability. We base the elliptical processes on a representation of elliptical distributions as a continuous mixture of Gaussian distributions and derive closed-form expressions for the marginal and conditional distributions. We perform numerical experiments on robust regression using an elliptical process defined by a piecewise constant mixing distribution, and show advantages compared with a Gaussian process. The elliptical processes may become a replacement for Gaussian processes in several settings, including when the likelihood is not Gaussian or when accurate tail modeling is critical.
△ Less
Submitted 2 December, 2020; v1 submitted 13 March, 2020;
originally announced March 2020.
-
A unified representation network for segmentation with missing modalities
Authors:
Kenneth Lau,
Jonas Adler,
Jens Sjölund
Abstract:
Over the last few years machine learning has demonstrated groundbreaking results in many areas of medical image analysis, including segmentation. A key assumption, however, is that the train- and test distributions match. We study a realistic scenario where this assumption is clearly violated, namely segmentation with missing input modalities. We describe two neural network approaches that can han…
▽ More
Over the last few years machine learning has demonstrated groundbreaking results in many areas of medical image analysis, including segmentation. A key assumption, however, is that the train- and test distributions match. We study a realistic scenario where this assumption is clearly violated, namely segmentation with missing input modalities. We describe two neural network approaches that can handle a variable number of input modalities. The first is modality dropout: a simple but surprisingly effective modification of the training. The second is the unified representation network: a network architecture that maps a variable number of input modalities into a unified representation that can be used for downstream tasks such as segmentation. We demonstrate that modality dropout makes a standard segmentation network reasonably robust to missing modalities, but that the same network works even better if trained on the unified representation.
△ Less
Submitted 19 August, 2019;
originally announced August 2019.
-
A linear programming approach to inverse planning in Gamma Knife radiosurgery
Authors:
Jens Sjölund,
Stella Riad,
Marcus Hennix,
Håkan Nordström
Abstract:
Leksell Gamma Knife is a stereotactic radiosurgery system that allows fine-grained control of the delivered dose distribution. We describe a new inverse planning approach that both resolves shortcomings of earlier approaches and unlocks new capabilities. We fix the isocenter positions and perform sector-duration optimization using linear programming, and study the effect of beam-on time penalizati…
▽ More
Leksell Gamma Knife is a stereotactic radiosurgery system that allows fine-grained control of the delivered dose distribution. We describe a new inverse planning approach that both resolves shortcomings of earlier approaches and unlocks new capabilities. We fix the isocenter positions and perform sector-duration optimization using linear programming, and study the effect of beam-on time penalization on the trade-off between beam-on time and plan quality. We also describe two techniques that reduce the problem size and thus further reduce the solution time: dualization and representative subsampling. The beam-on time penalization reduces the beam-on time by a factor 2-3 compared with the naive alternative. Dualization and representative subsampling each leads to optimization time-savings by a factor 5-20. Overall, we find in a comparison with 75 clinical plans that we can always find plans with similar coverage and better selectivity and beam-on time. In 44 of these, we can even find a plan that also has better gradient index. On a standard GammaPlan workstation, the optimization times ranged from 2.3 to 26 s with a median time of 5.7 s. In conclusion, we present a combination of techniques that enables sector-duration optimization in a clinically feasible time frame.
△ Less
Submitted 19 December, 2018; v1 submitted 8 June, 2018;
originally announced June 2018.
-
Bayesian uncertainty quantification in linear models for diffusion MRI
Authors:
Jens Sjölund,
Anders Eklund,
Evren Özarslan,
Magnus Herberthson,
Maria Bånkestad,
Hans Knutsson
Abstract:
Diffusion MRI (dMRI) is a valuable tool in the assessment of tissue microstructure. By fitting a model to the dMRI signal it is possible to derive various quantitative features. Several of the most popular dMRI signal models are expansions in an appropriately chosen basis, where the coefficients are determined using some variation of least-squares. However, such approaches lack any notion of uncer…
▽ More
Diffusion MRI (dMRI) is a valuable tool in the assessment of tissue microstructure. By fitting a model to the dMRI signal it is possible to derive various quantitative features. Several of the most popular dMRI signal models are expansions in an appropriately chosen basis, where the coefficients are determined using some variation of least-squares. However, such approaches lack any notion of uncertainty, which could be valuable in e.g. group analyses. In this work, we use a probabilistic interpretation of linear least-squares methods to recast popular dMRI models as Bayesian ones. This makes it possible to quantify the uncertainty of any derived quantity. In particular, for quantities that are affine functions of the coefficients, the posterior distribution can be expressed in closed-form. We simulated measurements from single- and double-tensor models where the correct values of several quantities are known, to validate that the theoretically derived quantiles agree with those observed empirically. We included results from residual bootstrap for comparison and found good agreement. The validation employed several different models: Diffusion Tensor Imaging (DTI), Mean Apparent Propagator MRI (MAP-MRI) and Constrained Spherical Deconvolution (CSD). We also used in vivo data to visualize maps of quantitative features and corresponding uncertainties, and to show how our approach can be used in a group analysis to downweight subjects with high uncertainty. In summary, we convert successful linear models for dMRI signal estimation to probabilistic models, capable of accurate uncertainty quantification.
△ Less
Submitted 19 February, 2018; v1 submitted 16 November, 2017;
originally announced November 2017.
-
Whole-brain diffusional variance decomposition (DIVIDE): Demonstration of technical feasibility at clinical MRI systems
Authors:
Filip Szczepankiewicz,
Jens Sjölund,
Freddy Ståhlberg,
Jimmy Lätt,
Markus Nilsson
Abstract:
Purpose: To assess the technical feasibility of whole-brain diffusional variance decomposition (DIVIDE) based on q-space trajectory encoding (QTE) at clinical MRI systems with varying performance. DIVIDE is used to separate diffusional heterogeneity into components that arise due to isotropic and anisotropic tissue structures. Methods: We designed imaging protocols for DIVIDE using numerically opt…
▽ More
Purpose: To assess the technical feasibility of whole-brain diffusional variance decomposition (DIVIDE) based on q-space trajectory encoding (QTE) at clinical MRI systems with varying performance. DIVIDE is used to separate diffusional heterogeneity into components that arise due to isotropic and anisotropic tissue structures. Methods: We designed imaging protocols for DIVIDE using numerically optimized gradient waveforms for diffusion encoding. Imaging was performed at systems with magnetic field strengths between 1.5 and 7 T, and gradient amplitudes between 33 and 80 mT/m. Technical feasibility was assessed from signal characteristics and quality of parameter maps in a single volunteer scanned at all systems. Results: The technical feasibility of QTE and DIVIDE was demonstrated at all systems. The system with the highest performance allowed whole-brain DIVIDE at 2 mm isotropic voxels. The system with the lowest performance required a spatial resolution of 2.5x2.5x4 mm3 to yield a sufficient signal-to-noise ratio. Conclusions: Whole-brain DIVIDE based on QTE is feasible at the investigated MRI systems. This demonstration indicates that tissue features beyond those accessible by conventional diffusion encoding may be explored on a wide range of MRI systems.
△ Less
Submitted 20 December, 2016;
originally announced December 2016.
-
Gaussian process regression can turn non-uniform and undersampled diffusion MRI data into diffusion spectrum imaging
Authors:
Jens Sjölund,
Anders Eklund,
Evren Özarslan,
Hans Knutsson
Abstract:
We propose to use Gaussian process regression to accurately estimate the diffusion MRI signal at arbitrary locations in q-space. By estimating the signal on a grid, we can do synthetic diffusion spectrum imaging: reconstructing the ensemble averaged propagator (EAP) by an inverse Fourier transform. We also propose an alternative reconstruction method guaranteeing a nonnegative EAP that integrates…
▽ More
We propose to use Gaussian process regression to accurately estimate the diffusion MRI signal at arbitrary locations in q-space. By estimating the signal on a grid, we can do synthetic diffusion spectrum imaging: reconstructing the ensemble averaged propagator (EAP) by an inverse Fourier transform. We also propose an alternative reconstruction method guaranteeing a nonnegative EAP that integrates to unity. The reconstruction is validated on data simulated from two Gaussians at various crossing angles. Moreover, we demonstrate on non-uniformly sampled in vivo data that the method is far superior to linear interpolation, and allows a drastic undersampling of the data with only a minor loss of accuracy. We envision the method as a potential replacement for standard diffusion spectrum imaging, in particular when acquistion time is limited.
△ Less
Submitted 9 November, 2016;
originally announced November 2016.
