Skip to content

Commit 996c019

Browse files
theo-barfoottheo-barfoot
authored
Merge pull request #72 update-paper
Updated JOSS Paper Draft
2 parents 3b1e478 + 9dc6d52 commit 996c019

4 files changed

Lines changed: 264 additions & 103 deletions

File tree

‎paper/jats/paper.jats‎

Lines changed: 190 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -57,10 +57,10 @@ PyTorch</article-title>
5757
</institution-wrap>
5858
</aff>
5959
</contrib-group>
60-
<pub-date date-type="pub" publication-format="electronic" iso-8601-date="2025-08-15">
61-
<day>15</day>
62-
<month>8</month>
63-
<year>2025</year>
60+
<pub-date date-type="pub" publication-format="electronic" iso-8601-date="2026-02-03">
61+
<day>3</day>
62+
<month>2</month>
63+
<year>2026</year>
6464
</pub-date>
6565
<volume>¿VOL?</volume>
6666
<issue>¿ISSUE?</issue>
@@ -90,21 +90,23 @@ a Creative Commons Attribution 4.0 International License (CC BY
9090
<sec id="summary">
9191
<title>Summary</title>
9292
<p>The <monospace>torchsparsegradutils</monospace> package provides
93-
gradient-preserving sparse tensor operations for PyTorch
93+
differentiable sparse linear-algebra utilities for PyTorch
9494
(<xref alt="Paszke et al., 2019" rid="ref-pytorch" ref-type="bibr">Paszke
95-
et al., 2019</xref>), addressing the critical limitation that
96-
PyTorch’s native sparse operations do not support sparse gradients
97-
during backpropagation. This package enables memory-efficient
98-
optimisation of high-dimensional models by maintaining sparsity
99-
patterns throughout the entire forward and backward pass computation,
100-
supporting both coordinate list (COO) and compressed sparse row (CSR)
101-
formats.</p>
102-
<p>Key features include: (1) memory-efficient sparse matrix
103-
multiplication with sparse gradient preservation, (2) sparse
104-
triangular and generic linear system solvers, enabling sparse
105-
gradients during backpropagation, and multiple algorithmic backends
106-
(BICGSTAB, CG, LSMR, MINRES), (3) cross-platform sparse solver
107-
wrappers for CuPy
95+
et al., 2019</xref>) that preserve sparsity for returned gradients
96+
during backpropagation. While PyTorch directly supports sparse
97+
tensors, its default semantics treat sparse layouts as storage
98+
optimisations rather than a mathematical structure that results in
99+
optimising directly for that sparse subspace. Gradients resulting from
100+
PyTorch native functions are often dense and incompatible with
101+
end-to-end training of models that require fixed sparsity patterns
102+
(e.g., sparse covariance/precision structures).</p>
103+
<p>To address this limitation, we introduce
104+
<monospace>torchsparsegradutils</monospace>. Key features include: (1)
105+
memory-efficient sparse-dense matrix multiplication with sparse
106+
gradient preservation, (2) sparse triangular and generic linear system
107+
solvers, enabling sparse gradients during backpropagation, and
108+
multiple algorithmic backends (BICGSTAB, CG, LSMR, MINRES), (3)
109+
cross-platform sparse solver wrappers for CuPy
108110
(<xref alt="Okuta et al., 2017" rid="ref-cupy" ref-type="bibr">Okuta
109111
et al., 2017</xref>) and JAX
110112
(<xref alt="Bradbury et al., 2018" rid="ref-jax" ref-type="bibr">Bradbury
@@ -117,41 +119,125 @@ a Creative Commons Attribution 4.0 International License (CC BY
117119
<mml:math display="inline" xmlns:mml="http://www.w3.org/1998/Math/MathML"><mml:mrow><mml:mi>𝐋</mml:mi><mml:mi>𝐃</mml:mi><mml:msup><mml:mi>𝐋</mml:mi><mml:mi>T</mml:mi></mml:msup></mml:mrow></mml:math></alternatives></inline-formula>
118120
sparse covariance and precision matrix parameterisations with
119121
reparameterised sampling methods, and (5) specialised encoders for
120-
spatial neighbourhood relationships in volumetric data.</p>
121-
<p>The package addresses PyTorch limitation of dense gradients
122-
resulting in memory errors, such as issue #41128, by implementing
123-
custom autograd functions that preserve sparsity in gradients,
124-
enabling practical training of models with sparse covariance and
125-
precision structures on high-dimensional data where dense alternatives
126-
become computationally intractable.</p>
122+
spatial neighbourhood relationships in N-dimensional data.</p>
127123
</sec>
128124
<sec id="statement-of-need">
129125
<title>Statement of need</title>
130-
<p>Sparse tensors are essential for computationally tractable machine
131-
learning on high-dimensional data, yet PyTorch’s sparse tensor
132-
operations suffer from a critical limitation: gradients are computed
133-
in dense format even when the forward pass maintains sparsity. This
134-
results in prohibitive memory usage that scales quadratically with
135-
problem dimension rather than linearly with the number of non-zero
136-
elements (nnz).</p>
137-
<p>For applications requiring sparse covariance modelling—such as
138-
medical imaging with millions of voxels, spatial statistics, and
139-
large-scale Gaussian processes—dense gradient computation renders
140-
optimisation infeasible. A sparse covariance matrix with 1 million
141-
dimensions and 0.1% sparsity contains 1 billion non-zero elements, but
142-
dense gradients would require storing 1 trillion parameters.</p>
143-
<p>This package solves this fundamental limitation by implementing
144-
custom autograd functions that preserve sparsity patterns throughout
145-
both forward and backward passes. Our sparse multivariate normal
146-
distribution enables optimisation of million-dimensional Gaussian
147-
models with memory usage scaling as O(nnz) rather than O(n²), where n
148-
is the dimension of the multivariate distribution.</p>
149-
<p>Beyond memory efficiency, the package addresses the algorithmic
150-
challenge of sparse linear system solving by providing multiple
151-
iterative solver backends with automatic differentiation support. This
152-
enables end-to-end optimisation of complex probabilistic models that
153-
would be computationally intractable with existing PyTorch sparse
154-
operations.</p>
126+
<p>Many scientific machine learning models benefit from representing
127+
large linear operators (e.g., neighbourhood couplings, precision
128+
matrices, sparse Jacobians) using sparse tensors to reduce memory and
129+
compute complexity. In high-dimensional settings, such as volumetric
130+
medical imaging, dense covariance or precision parameterisations are
131+
typically intractable, motivating sparse end-to-end
132+
parameterisations.</p>
133+
<p>However, learning these models with gradient-based optimisation
134+
requires backpropagation through sparse linear algebra (matrix
135+
products, triangular solves, and linear system solves). PyTorch’s
136+
default sparse semantics are not designed to preserve user-imposed
137+
sparsity structure during differentiation (PyTorch issue #87448),
138+
which can lead to memory blow-ups and prevent end-to-end optimisation
139+
of sparse probabilistic models.</p>
140+
<p><monospace>torchsparsegradutils</monospace> addresses this gap by
141+
implementing custom autograd functions for key sparse operators that
142+
return gradients only for stored nonzeros, enabling practical
143+
optimisation of models that rely on fixed sparse structure, such as
144+
sparse multivariate normal distributions with sparse
145+
covariance/precision factors.</p>
146+
</sec>
147+
<sec id="state-of-the-field">
148+
<title>State of the field</title>
149+
<p>PyTorch
150+
(<xref alt="Paszke et al., 2019" rid="ref-pytorch" ref-type="bibr">Paszke
151+
et al., 2019</xref>) exposes sparse layouts (COO, CSR, and related
152+
formats) and implements a growing set of sparse operations. However,
153+
PyTorch’s design goal is <italic>dense-equivalent semantics</italic>
154+
for sparse layouts: a guiding invariant is that applying an operation
155+
in sparse form should match applying it in dense form after
156+
conversion, including the backward function (PyTorch issue #87448).
157+
This makes it difficult to learn parameters that are intended to
158+
remain structurally sparse, because gradients may be produced for
159+
implicit zeros, or intermediate computations may densify.</p>
160+
<p>PyTorch also provides <monospace>MaskedTensor</monospace>,
161+
distringuishing specified and unspecified elements in tensors and is
162+
conceptually closer to the constrained-subspace interpretation of
163+
sparsity. However, <monospace>MaskedTensor</monospace> remains at
164+
prototype stage with incomplete operator coverage, and storing a full
165+
boolean mask incurs a significant memory overhead, partially negating
166+
the memory benefits of sparse index-based representations for
167+
large-scale problems.</p>
168+
<p>Other libraries provide efficient sparse kernels but do not
169+
directly solve “sparsity-preserving gradients in PyTorch”: SciPy
170+
(<xref alt="Virtanen et al., 2020" rid="ref-scipy" ref-type="bibr">Virtanen
171+
et al., 2020</xref>) provides mature sparse linear algebra but no
172+
automatic differentiation; CuPy
173+
(<xref alt="Okuta et al., 2017" rid="ref-cupy" ref-type="bibr">Okuta
174+
et al., 2017</xref>) and JAX
175+
(<xref alt="Bradbury et al., 2018" rid="ref-jax" ref-type="bibr">Bradbury
176+
et al., 2018</xref>) provide sparse solvers in their respective
177+
ecosystems but are not drop-in components for PyTorch
178+
autograd/training loops. GPyTorch
179+
(<xref alt="Gardner et al., 2018" rid="ref-gpytorch" ref-type="bibr">Gardner
180+
et al., 2018</xref>) targets scalable Gaussian process inference via
181+
kernel structure and approximations (e.g., inducing/structured
182+
methods) rather than arbitrary user-specified sparse
183+
covariance/precision factors. PyTorch Geometric’s torch_sparse
184+
(<xref alt="Fey &amp; Lenssen, 2019" rid="ref-pytorch_geometric" ref-type="bibr">Fey
185+
&amp; Lenssen, 2019</xref>) focuses on graph message-passing
186+
primitives rather than sparse covariance/precision modelling and
187+
differentiable sparse solves for probabilistic models.</p>
188+
</sec>
189+
<sec id="software-design">
190+
<title>Software design</title>
191+
<p><monospace>torchsparsegradutils</monospace> is built around
192+
<monospace>torch.autograd.Function</monospace> operators that wrap
193+
PyTorch’s forward sparse kernels but override the backward pass to
194+
preserve sparsity for selected inputs. This design keeps the
195+
user-facing API close to standard PyTorch code while making sparsity
196+
preservation an explicit, opt-in choice.</p>
197+
<p>Two design trade-offs shaped the implementation. First, the package
198+
targets <italic>structure-preserving learning</italic> over maximal
199+
operator coverage, as only a focused set of operations (sparse matrix
200+
products, triangular solves, generic sparse solvers) are implemented,
201+
but these are sufficient to support sparse multivariate normal
202+
sampling and sparse solver-based models. Second, for broad
203+
device/backend compatibility, the package combines native PyTorch
204+
implementations (iterative Krylov solvers: CG, BiCGSTAB, LSMR, MINRES)
205+
with optional wrappers to external libraries (CuPy, JAX), allowing
206+
users to trade off portability versus performance.</p>
207+
<p><bold>Build vs. contribute justification.</bold> PyTorch’s current
208+
semantics treat sparse layouts as performance optimisations and
209+
prioritise the dense-equivalence invariant (PyTorch issue #87448). In
210+
contrast, this package intentionally provides
211+
<italic>structure-preserving</italic> backward passes for specific
212+
operators to enable learning with fixed sparsity patterns (e.g.,
213+
sparse triangular factors for covariance/precision). This difference
214+
is semantic (not just implementation), so the functionality is better
215+
delivered as an opt-in external library rather than changing PyTorch’s
216+
default behaviour.</p>
217+
</sec>
218+
<sec id="research-impact-statement">
219+
<title>Research impact statement</title>
220+
<p>This software provides an opt-in path to sparsity-preserving
221+
gradients for sparse linear algebra in PyTorch, enabling research
222+
prototypes that would otherwise be limited by dense gradients or
223+
densification. The package is currently being used in active research
224+
projects for medical image segmentation, though publications resulting
225+
from this work are still in preparation.</p>
226+
<p>The codebase demonstrates community-readiness through comprehensive
227+
infrastructure: documentation with quickstart guides and API
228+
references, extensive test coverage across all modules, CI/CD
229+
pipelines for automated testing, and an open contribution process via
230+
GitHub issues and pull requests. The codebase has been developed
231+
openly over multiple years with public commit history, releases, and
232+
issue tracking. Benchmark suites comparing solver performance across
233+
problem sizes and sparsity patterns provide reproducible reference
234+
materials.</p>
235+
<p>Given the broad applicability of sparse structured
236+
Gaussians—spanning medical imaging, spatial statistics, geostatistics,
237+
and large-scale probabilistic modelling, we anticipate growing
238+
adoption as the research community increasingly requires
239+
memory-efficient optimisation of high-dimensional probabilistic
240+
models.</p>
155241
</sec>
156242
<sec id="mathematics">
157243
<title>Mathematics</title>
@@ -288,19 +374,24 @@ a Creative Commons Attribution 4.0 International License (CC BY
288374
matrices by avoiding strict positive definiteness constraints.</p>
289375
</sec>
290376
</sec>
377+
<sec id="ai-usage-disclosure">
378+
<title>AI usage disclosure</title>
379+
<p>Generative AI tools were used during development of this software
380+
and manuscript. Various large language models were used to assist with
381+
code generation, refactoring, and test scaffolding for portions of the
382+
codebase, and AI assistance was used to draft and edit parts of the
383+
documentation and this manuscript. The repository was initiated prior
384+
to widespread AI coding assistant adoption, with AI tools incorporated
385+
during later development phases. All AI-assisted outputs were
386+
reviewed, edited, and validated by the human authors, who take
387+
responsibility for the final software and paper.</p>
388+
</sec>
291389
<sec id="acknowledgements">
292390
<title>Acknowledgements</title>
293-
<p>The authors acknowledge the PyTorch development team for providing
294-
the foundational sparse tensor infrastructure. We thank the SciPy
295-
(<xref alt="Virtanen et al., 2020" rid="ref-scipy" ref-type="bibr">Virtanen
296-
et al., 2020</xref>), CuPy
297-
(<xref alt="Okuta et al., 2017" rid="ref-cupy" ref-type="bibr">Okuta
298-
et al., 2017</xref>), and JAX
299-
(<xref alt="Bradbury et al., 2018" rid="ref-jax" ref-type="bibr">Bradbury
300-
et al., 2018</xref>) communities for high-performance sparse linear
301-
algebra implementations. Algorithm implementations adapt and extend
302-
methods from pykrylov (BICGSTAB), cornellius-gp/linear_operator (CG,
303-
MINRES), and pytorch-minimize (LSMR)
391+
<p>We thank the PyTorch development team for foundational sparse
392+
tensor support. We also acknowledge upstream solver implementations
393+
and references used as starting points for iterative methods
394+
(pykrylov, cornellius-gp/linear_operator, pytorch-minimize)
304395
(<xref alt="Saad, 2003" rid="ref-saad2003iterative" ref-type="bibr">Saad,
305396
2003</xref>). We thank Floris Laporte for his excellent tutorial on
306397
implementing sparse linear system solvers in PyTorch
@@ -334,6 +425,7 @@ a Creative Commons Attribution 4.0 International License (CC BY
334425
<source>Advances in neural information processing systems</source>
335426
<year iso-8601-date="2019">2019</year>
336427
<volume>32</volume>
428+
<uri>https://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library</uri>
337429
<fpage>8024</fpage>
338430
<lpage>8035</lpage>
339431
</element-citation>
@@ -348,9 +440,11 @@ a Creative Commons Attribution 4.0 International License (CC BY
348440
<name><surname>Loomis</surname><given-names>Crissman</given-names></name>
349441
</person-group>
350442
<article-title>CuPy: A NumPy-compatible library for NVIDIA GPU calculations</article-title>
351-
<source>Proceedings of the workshop on machine learning systems (LearningSys) at NeurIPS</source>
443+
<source>Proceedings of the workshop on machine learning systems (LearningSys) at the 31st conference on neural information processing systems</source>
352444
<year iso-8601-date="2017">2017</year>
353445
<uri>https://learningsys.org/nips17/assets/papers/paper_16.pdf</uri>
446+
<fpage>1</fpage>
447+
<lpage>7</lpage>
354448
</element-citation>
355449
</ref>
356450
<ref id="ref-jax">
@@ -388,7 +482,7 @@ a Creative Commons Attribution 4.0 International License (CC BY
388482
<name><surname>Bright</surname><given-names>Jonathan</given-names></name>
389483
<name><surname>others</surname></name>
390484
</person-group>
391-
<article-title>SciPy 1.0: Fundamental algorithms for scientific computing in python</article-title>
485+
<article-title>SciPy 1.0: Fundamental algorithms for scientific computing in Python</article-title>
392486
<source>Nature Methods</source>
393487
<year iso-8601-date="2020">2020</year>
394488
<volume>17</volume>
@@ -404,10 +498,40 @@ a Creative Commons Attribution 4.0 International License (CC BY
404498
<name><surname>Saad</surname><given-names>Yousef</given-names></name>
405499
</person-group>
406500
<source>Iterative methods for sparse linear systems</source>
407-
<publisher-name>SIAM</publisher-name>
501+
<publisher-name>Society for Industrial; Applied Mathematics</publisher-name>
408502
<year iso-8601-date="2003">2003</year>
409-
<edition>2nd</edition>
410-
<uri>https://www-users.cse.umn.edu/~saad/IterMethBook_2ndEd.pdf</uri>
503+
<edition>2</edition>
504+
<uri>https://epubs.siam.org/doi/10.1137/1.9780898718003</uri>
505+
<pub-id pub-id-type="doi">10.1137/1.9780898718003</pub-id>
506+
</element-citation>
507+
</ref>
508+
<ref id="ref-pytorch_geometric">
509+
<element-citation publication-type="paper-conference">
510+
<person-group person-group-type="author">
511+
<name><surname>Fey</surname><given-names>Matthias</given-names></name>
512+
<name><surname>Lenssen</surname><given-names>Jan E.</given-names></name>
513+
</person-group>
514+
<article-title>Fast graph representation learning with PyTorch Geometric</article-title>
515+
<source>ICLR workshop on representation learning on graphs and manifolds</source>
516+
<year iso-8601-date="2019">2019</year>
517+
<uri>https://arxiv.org/abs/1903.02428</uri>
518+
</element-citation>
519+
</ref>
520+
<ref id="ref-gpytorch">
521+
<element-citation publication-type="paper-conference">
522+
<person-group person-group-type="author">
523+
<name><surname>Gardner</surname><given-names>Jacob R.</given-names></name>
524+
<name><surname>Pleiss</surname><given-names>Geoff</given-names></name>
525+
<name><surname>Bindel</surname><given-names>David</given-names></name>
526+
<name><surname>Weinberger</surname><given-names>Kilian Q.</given-names></name>
527+
<name><surname>Wilson</surname><given-names>Andrew Gordon</given-names></name>
528+
</person-group>
529+
<article-title>GPyTorch: Blackbox matrix-matrix gaussian process inference with GPU acceleration</article-title>
530+
<source>Advances in neural information processing systems</source>
531+
<year iso-8601-date="2018">2018</year>
532+
<volume>31</volume>
533+
<uri>https://arxiv.org/abs/1809.11165</uri>
534+
<pub-id pub-id-type="doi">10.5555/3327757.3327857</pub-id>
411535
</element-citation>
412536
</ref>
413537
<ref id="ref-flaport2020sparse">

0 commit comments

Comments
 (0)