Skip to content

RESEARCH_OPPORTUNITIES.md

Latest commit

 

History

History
76 lines (50 loc) · 8.88 KB

RESEARCH_OPPORTUNITIES.md

File metadata and controls

76 lines (50 loc) · 8.88 KB
Raw

Further Research Opportunities for mLR

This document outlines open research directions for the Memoization-based Scalable Laminography Reconstruction (mLR) project, grounded in the current implementation's design choices and observed limitations.

1. Adaptive Similarity Threshold

The current system uses a fixed similarity threshold (τ = 0.85–0.90) for the entire reconstruction. The rebuttal data shows a clear accuracy–speedup trade-off: τ=0.90 yields 90.1% accuracy at 66.2% normalized time, while τ=0.94 reaches 95.8% accuracy but only saves 12.6% of time.

Opportunities:

  • Iteration-aware threshold scheduling. Early ADMM iterations tolerate more approximation error since the solution is far from convergence. A decaying threshold (e.g., τ starts at 0.85 and anneals toward 0.96 in later outer iterations) could accelerate early iterations without sacrificing final reconstruction quality.
  • Spatially-varying thresholds. Regions of the reconstruction volume with low gradient magnitude or high similarity across iterations could tolerate more aggressive caching, while high-frequency regions (edges, fine features) demand recomputation. Per-chunk threshold adaptation based on local signal characteristics is unexplored.
  • Convergence-guided caching. Use ADMM primal/dual residuals (already computed in update_penalty) to dynamically decide when to relax or tighten the threshold rather than relying on a fixed schedule.

2. Encoder Architecture and Training

The current Complex2vec encoder is a shallow CNN (2 conv layers + pooling + FC) with a hardcoded linear layer size (55,296 input features). In practice, the encoding path in fftcl.py currently bypasses the CNN and uses AdaptiveAvgPool2d((3, 3)) followed by flattening — a much simpler pooling-based encoding.

Opportunities:

  • Lightweight encoding alternatives. Investigate whether structured hashing (e.g., locality-sensitive hashing) or random projections can replace the CNN encoder entirely while maintaining similar cache hit rates. This would eliminate the ~200 min training cost and remove the dependency on pre-trained model files.
  • Self-supervised encoder training. The current RegressionLoss trains the encoder to predict Euclidean distance between chunk pairs. Contrastive learning objectives (e.g., SimCLR-style loss) could produce embeddings with better separation in the similarity space, improving cache hit/miss discrimination near the threshold boundary.
  • Adjoint-direction encoding. The encoder currently only handles the forward direction (direction='fwd'); the adjoint path prints a placeholder message. Extending memoization to the adjoint operator would double the potential cache hits per ADMM iteration.
  • Input-adaptive embedding dimension. The embedding dimensions (dim1=96, dim2=144) are fixed. Investigating the relationship between embedding dimension, FAISS search accuracy, and cache effectiveness could reveal a more optimal operating point.

3. Multi-GPU Scaling and Communication

Latency measurements show that scaling from 1 GPU (~0.5 ms baseline) to 16 GPUs introduces outlier latencies of 1.8–4.5 seconds. The current multi-GPU strategy uses Python threading with per-GPU CUDA streams and relies on MPI for inter-node communication.

Opportunities:

  • Communication-computation overlap. The current chunked pipeline synchronizes all streams at each iteration step (stream1.synchronize(); stream2.synchronize(); stream3.synchronize()). Relaxing these synchronization barriers — for instance, allowing the next chunk's CPU→GPU transfer to overlap with the current chunk's computation — could hide transfer latency. This is partially implemented in the double-buffering pattern but not fully exploited.
  • NCCL-based collective operations. Replacing MPI point-to-point communication with NCCL collectives for GPU-to-GPU data movement would bypass CPU staging and reduce latency for multi-GPU gradient aggregation and scatter/gather operations.
  • Topology-aware work partitioning. The current equal-partition strategy (n1 // nworkers) does not account for NUMA topology or NVLink connectivity. Non-uniform partitioning that assigns more work to GPUs with faster interconnects could reduce tail latency.
  • Profiling and eliminating synchronization outliers. The large latency spikes (0.5 s at 4 GPUs, 4.5 s at 16 GPUs) in the benchmark data suggest synchronization stalls or memory allocation on the critical path. Systematic profiling with NVIDIA Nsight Systems could identify whether these are caused by page faults, CUDA context switches, or MPI barrier delays.

4. Distributed Memoization System

The DistributedFaissManager sends key-value pairs over MPI to a remote node that maintains both a FAISS index and a Redis metadata store. Several aspects of this system are currently constrained.

Opportunities:

  • Eliminating Redis dependency. The metadata store uses Redis purely as a key-value store for mapping FAISS index IDs to cached output arrays. An in-memory dictionary (as used locally in FaissVectorDB.vector_values) would remove the Redis deployment requirement and serialization overhead. Redis would only become justified if persistence across job restarts or cross-job sharing is needed.
  • Hierarchical caching (L1/L2). Implement a two-level cache: a fast, small local GPU-memory cache (L1) for the most recent chunks, backed by a larger CPU-memory or distributed cache (L2). This mirrors hardware cache hierarchies and could reduce MPI round trips for high-locality access patterns.
  • Asynchronous cache population. Currently, cache miss blocks are computed, then stored synchronously. Decoupling storage from the computation pipeline (e.g., storing results in a background thread after the output is already consumed) would remove storage from the critical path.
  • Cache eviction and memory management. The current cache grows without bound — vector_values accumulates entries indefinitely. Implementing LRU or frequency-based eviction would bound memory usage, particularly important for long-running reconstructions with many ADMM iterations.
  • Shared cross-iteration caching. The similarity profiling (test_similarity.py) shows that chunks at the same spatial position maintain high similarity across ADMM iterations. A persistent cache that spans outer iterations (rather than per-iteration caching) could amortize the cold-start cost.

5. Memory Management and Offloading

The admm_offload and admm_offload_async variants save/load intermediate arrays (psi, lamd, h) to /local/scratch/ using np.save/np.load. This is a straightforward but I/O-bound approach.

Opportunities:

  • Memory-mapped arrays. Replace explicit save/load cycles with np.memmap backed by NVMe storage, allowing the OS page cache to manage residency. This would eliminate explicit I/O calls and enable lazy loading.
  • Unified memory with prefetching. CUDA unified memory (cupy.cuda.ManagedMemory) with explicit prefetch hints could automate CPU↔GPU data movement for the offload variants, replacing the manual pinned memory + stream copy pipeline.
  • Compression of offloaded data. The intermediate arrays (psi, lamd, h) are complex64. Lossy compression (e.g., quantizing to float16 or using structured sparsity) before offloading could reduce I/O bandwidth requirements with minimal impact on convergence.
  • Pipelined I/O. The async variant (admm_offload_async) uses asyncio but still contains several await calls that serialize I/O. Restructuring to use a producer-consumer pattern with pre-fetching of next-iteration data during current-iteration computation would better utilize I/O bandwidth.

6. Algorithmic Extensions

Opportunities:

  • Beyond ADMM. The ADMM solver uses a fixed splitting with TV regularization. Investigating Primal-Dual Hybrid Gradient (PDHG) or linearized ADMM could yield faster convergence, especially when combined with memoization (fewer iterations means fewer total FFT evaluations).
  • Learned regularization. Replace hand-crafted TV regularization with a learned denoiser prior (Plug-and-Play or RED framework). This would require the denoiser to operate efficiently on the chunk-level data layout but could significantly improve reconstruction quality per iteration.
  • Stochastic/mini-batch updates. Rather than using all projection angles per iteration, stochastic ADMM with random angle subsets could reduce per-iteration cost. Combined with memoization, previously computed projections for excluded angles could be reused from cache.
  • Joint encoder-reconstruction optimization. Currently the encoder is trained offline on pre-collected data. End-to-end fine-tuning of the encoder during reconstruction (treating cache hit rate as a differentiable objective) could adapt the encoder to the specific dataset being reconstructed.