Document configuration of sparsity pattern representations (#256)

* Update README

* Update citation

* Add doc page on performance

* Update CHANGELOG

Author: adrhill

CHANGELOG.md

@@ -4,8 +4,8 @@
 
 This release is only breaking for users who touched unexported internals.
 
-* ![BREAKING][badge-breaking] Refactor internal types ([#254]).
-* ![Feature][badge-feature] Select internal data structures for sparsity pattern representations via `TracerSparsityDetector` and `TracerLocalSparsityDetector` keyword arguments `gradient_pattern_type` and `hessian_pattern_type` ([#254]).
+* ![BREAKING][badge-breaking] Refactor internal types ([#254])
+* ![Feature][badge-feature] Select internal data structures for sparsity pattern representations via the `TracerSparsityDetector` and `TracerLocalSparsityDetector` keyword arguments `gradient_pattern_type` and `hessian_pattern_type`. Refer to the new doc page on "Performance" for more information. ([#254], [#256])
 * ![Deprecation][badge-deprecation] Move DataInterpolations.jl extension out of SCT ([#255])
 
 ## Version `v0.6.21`
@@ -136,6 +136,7 @@ This release is only breaking for users who touched unexported internals.
 [badge-maintenance]: https://img.shields.io/badge/maintenance-gray.svg
 [badge-docs]: https://img.shields.io/badge/docs-orange.svg
 
+[#256]: https://github.com/adrhill/SparseConnectivityTracer.jl/pull/256
 [#255]: https://github.com/adrhill/SparseConnectivityTracer.jl/pull/255
 [#254]: https://github.com/adrhill/SparseConnectivityTracer.jl/pull/254
 [#249]: https://github.com/adrhill/SparseConnectivityTracer.jl/pull/249

CITATION.cff

@@ -12,18 +12,16 @@ version: 0.1.0
 date-released: 2024-04-22
 url: "https://github.com/adrhill/SparseConnectivityTracer.jl"
 preferred-citation:
-  type: generic
+  type: article
   authors:
   - family-names: "Hill"
     given-names: "Adrian"
     orcid: "https://orcid.org/0009-0009-5977-301X"
   - family-names: "Dalle"
     given-names: "Guillaume"
     orcid: "https://orcid.org/0000-0003-4866-1687"
-  title: "Sparser, Better, Faster, Stronger: Efficient Automatic Differentiation for Sparse Jacobians and Hessians"
+  title: "Sparser, Better, Faster, Stronger: Sparsity Detection for Efficient Automatic Differentiation"
+  journal: "Transactions on Machine Learning Research"
   year: 2025
-  eprint: "2501.17737"
-  archivePrefix: "arXiv"
-  primaryClass: "cs.LG"
-  url: "https://arxiv.org/abs/2501.17737"
-  doi: "10.48550/arXiv.2501.17737"
+  issn: "2835-8856"
+  url: "https://openreview.net/forum?id=GtXSN52nIW"

README.md

@@ -66,6 +66,13 @@ julia> jacobian_sparsity(layer, x, detector)
 ⎣⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢿⣦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠛⢿⣦⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠻⣷⣄⎦
 ```
 
+By default, `BitSet` is used for internal sparsity pattern representations.
+For very large inputs, it might be more efficient to set the type to `Set{UInt}`:
+
+```julia-repl
+julia> detector = TracerSparsityDetector(; gradient_pattern_type=Set{UInt})
+```
+
 ### Hessian
 
 For scalar functions `y = f(x)`, the sparsity pattern of the Hessian of $f$ can be obtained
@@ -95,6 +102,14 @@ julia> hessian_sparsity(g, x, detector)
  ⋅  1  ⋅  ⋅  1
 ```
 
+
+By default, a dictionaries of `BitSet` are used for internal sparsity pattern representations.
+For very large inputs, it might be more efficient to set the type to `Dict{UInt, Set{UInt}}`:
+
+```julia-repl
+julia> detector = TracerSparsityDetector(; hessian_pattern_type=Dict{UInt, Set{UInt}})
+```
+
 For more detailed examples, take a look at the [documentation](https://adrianhill.de/SparseConnectivityTracer.jl/stable).
 
 ### Local tracing
@@ -141,14 +156,14 @@ In fact, the functions `jacobian_sparsity` and `hessian_sparsity` are re-exporte
 If you use SparseConnectivityTracer in your research, please cite our preprint [*Sparser, Better, Faster, Stronger: Efficient Automatic Differentiation for Sparse Jacobians and Hessians*](https://arxiv.org/abs/2501.17737):
 
 ```bibtex
-@misc{hill2025sparserbetterfasterstronger,
-      title={Sparser, Better, Faster, Stronger: Efficient Automatic Differentiation for Sparse Jacobians and Hessians}, 
-      author={Adrian Hill and Guillaume Dalle},
-      year={2025},
-      eprint={2501.17737},
-      archivePrefix={arXiv},
-      primaryClass={cs.LG},
-      url={https://arxiv.org/abs/2501.17737}, 
+@article{hill2025sparser,
+  title={Sparser, Better, Faster, Stronger: Sparsity Detection for Efficient Automatic Differentiation},
+  author={Adrian Hill and Guillaume Dalle},
+  journal={Transactions on Machine Learning Research},
+  issn={2835-8856},
+  year={2025},
+  url={https://openreview.net/forum?id=GtXSN52nIW},
+  note={}
 }
 ```
 

docs/Project.toml

@@ -1,7 +1,9 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterMermaid = "a078cd44-4d9c-4618-b545-3ab9d77f9177"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 

docs/make.jl

@@ -25,7 +25,7 @@ makedocs(;
     pages = [
         "Getting Started" => "index.md",
         "User Documentation" =>
-            ["user/global_vs_local.md", "user/limitations.md", "user/api.md"],
+            ["user/global_vs_local.md", "user/performance.md", "user/limitations.md", "user/api.md"],
         "Developer Documentation" => [
             "internals/how_it_works.md",
             "internals/adding_overloads.md",

docs/src/user/performance.md

@@ -0,0 +1,48 @@
+# [Performance](@id performance)
+
+## Data structures for sparsity pattern representations
+
+The most efficient internal data structure for sparsity pattern representations
+depends on the number of inputs and the computational graph / sparsity of a given function.
+
+Let's use a convolutional layer from Flux.jl as an example.
+By default, SCT uses `BitSet` for Jacobian sparsity detection, which is well suited for small to medium sized functions.
+
+```@example Flux
+using SparseConnectivityTracer, Flux, BenchmarkTools
+
+x = rand(28, 28, 3, 1)
+layer = Conv((3, 3), 3 => 2)
+
+detector_bitset = TracerSparsityDetector()
+jacobian_sparsity(layer, x, detector_bitset)
+```
+
+```@example Flux
+@benchmark jacobian_sparsity(layer, x, detector_bitset);
+```
+
+Instead of `BitSet`, we can use any concrete subtype of `AbstractSet{<:Integer}`, for example `Set{UInt}`.
+To set the sparsity pattern type for Jacobian sparsity detection, we use the keyword argument `gradient_pattern_type`:
+
+```@example Flux
+detector_set = TracerSparsityDetector(; gradient_pattern_type=Set{UInt})
+@benchmark jacobian_sparsity(layer, x, detector_set);
+```
+
+While this is slower for the given input size, the performance is highly dependant on the problem.
+For larger inputs (e.g. of size $224 \times 224 \times 3 \times 1$), `detector_set` will outperform `detector_bitset`.
+Note that memory requirement will vary as well.
+
+For Hessians sparsity detection, the internal sparsity pattern representation uses either concrete subtypes of
+`AbstractDict{I, AbstractSet{I}}` or `AbstractSet{Tuple{I, I}}`, where `I <: Integer`.
+By default, `Dict{Int, BitSet)` is used.
+To set the sparsity pattern type, use the keyword argument `hessian_pattern_type`:
+
+```@example Flux
+detector = TracerSparsityDetector(; hessian_pattern_type=Dict{UInt, Set{UInt}})
+```
+
+Data structures can also be set analogously for `TracerLocalSparsityDetector`.
+If both Jacobian and Hessian sparsity patterns are needed, 
+`gradient_pattern_type` and `hessian_pattern_type` can be set separately.

src/adtypes_interface.jl

@@ -36,7 +36,7 @@ const DOC_KWARGS = """# Keyword arguments
   Defaults to `$DEFAULT_SET_TYPE`.
 - `hessian_pattern_type::Type`: 
   Data structure used for bookkeeping of Hessian sparsity patters, used in `hessian_sparsity`.
-  Supports concrete subtypes of `AbstractDict{I, AbstractSet{I}}` or `AbstractSet{Tuple{I, I}}}`, where `I <: Integer`.
+  Supports concrete subtypes of `AbstractDict{I, AbstractSet{I}}` or `AbstractSet{Tuple{I, I}}`, where `I <: Integer`.
   Defaults to `$DEFAULT_DICT_TYPE`.
 - `shared_hessian_pattern::Bool`: 
   Indicate whether second-order information in Hessian sparsity patterns **always** shares memory and whether operators are **allowed** to mutate `HessianTracers`.