Upload folder using huggingface_hub

paper/uraionspec-paper.pdf

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ad6156cd8842d4021aa11427a65bfacfa05f46dac101fa1efbe068d23f3cf164
-size 490909
+oid sha256:f93345bf1b491a8d754efba0c5146698a9a8262b8961ed1151af31290b48250f
+size 500974

paper/uraionspec-paper.tex

@@ -16,6 +16,11 @@
 % enumitem not available — using basic LaTeX lists instead
 % caption/subcaption not available — using basic figure handling
 
+% Fix hyperref destination conflict with algorithm float
+\makeatletter
+\providecommand*{\theHalgorithm}{\thealgorithm}
+\makeatother
+
 \geometry{margin=1in}
 
 % ---------- metadata for arXiv ----------
@@ -67,13 +72,14 @@ backbone with a lightweight sequential head, and confidence-scheduled
 verification that dynamically adjusts verification length based on predicted
 acceptance probabilities and engine throughput.
 
-We present \UraionSpec, a clean, modular, and verified implementation of the
-\DSpark algorithm. Our implementation faithfully reproduces the core
+We present \UraionSpec, a clean, modular, verified implementation of
+\DSpark. It faithfully reproduces the core
 components---the Markov and RNN sequential heads, confidence head, Sequential
-Temperature Scaling (STS) calibration, hardware-aware prefix scheduler
-(Algorithm~1), and the three-term training objective---while adding a
-\DFlash-style parallel backbone with target model KV injection that was
-missing from the initial release. With 80 passing unit tests, a fully
+Temperature Scaling~(STS) calibration, hardware-aware prefix scheduler
+(Algorithm~1), and the three-term training objective---and includes a
+\DFlash-style parallel backbone with target model KV injection as a
+standalone module for integration into the draft model pipeline.
+With 80 passing unit tests, a fully
 configurable training pipeline, comprehensive evaluation tools, and
 end-to-end smoke tests, \UraionSpec is designed for reproducible research
 and practical experimentation. The codebase is publicly available on
@@ -109,7 +115,7 @@ contributions:
     (the \DFlash architecture) processes all proposal positions in a
     single forward pass, while a lightweight sequential head (Markov or RNN)
     injects inter-token dependency. This combines the speed of parallel
-    drafters~\cite{stern2018blockwise, xia2022medusa} with the quality of
+    drafters~\cite{stern2018blockwise, xia2024medusa} with the quality of
     autoregressive ones~\cite{miao2025eagle3}.
   \item \textbf{Confidence-scheduled verification}: A trained confidence head
     predicts per-position acceptance probabilities, and a hardware-aware
@@ -121,15 +127,15 @@ contributions:
 \paragraph{Our contribution.}
 We present \UraionSpec, a faithful, modular, and verified implementation of
 \DSpark. Our implementation includes all core components listed in the
-original paper, plus a \DFlash-style parallel backbone with target model KV
-injection that was identified as a critical gap in initial implementations.
+original paper, and a standalone \DFlash-style parallel backbone with target model KV
+injection module, available for future integration into the draft model pipeline.
 The codebase is designed for clarity and reproducibility rather than
 production scale, making it suitable for academic research, ablation studies,
 and educational purposes. Key features include:
 
 \begin{itemize}
   \item Complete implementation of all \DSpark components (80 unit tests)
-  \item \DFlash-style backbone with target model KV injection
+  \item \DFlash-style backbone module with target model KV injection
   \item Markov, gated Markov, and RNN sequential heads
   \item Confidence head with analytical acceptance-rate supervision
   \item Sequential Temperature Scaling (STS) calibration
@@ -156,7 +162,7 @@ is accepted with probability
   \label{eq:acceptance}
 \end{equation}
 and the first rejection marks the end of the accepted prefix. A bonus token
-is then sampled from the target's residual distribution at the rejection
+is then sampled from the target model's residual distribution at the rejection
 position, ensuring the final distribution matches the target model exactly.
 
 The expected number of accepted tokens $\tau$ is
@@ -234,7 +240,7 @@ UraionSpec/
   src/uraionspec/
     models/          -- Draft model components
       markov_head.py       -- Low-rank transition bias
-      rnn_head.py          -- GRU-like recurrent head
+      rnn_head.py          -- GRU (Gated Recurrent Unit)-like recurrent head
       confidence_head.py   -- Per-position acceptance predictor
       dflash_backbone.py   -- DFlash backbone w/ KV injection
       draft_model.py       -- Combined DSpark draft model
@@ -286,7 +292,7 @@ This allows the backbone representation to influence how much Markov bias to
 apply, providing context-dependent transition modeling.
 
 \paragraph{RNN head.}
-For full prefix history, we implement a GRU-like recurrent head:
+For full prefix history, we implement a GRU (Gated Recurrent Unit)-like recurrent head:
 \begin{align}
   z_k &= [s_{k-1} ;\; W_1[x_{k-1}] ;\; h_k], \\
   s_k &= \sigma(W_g z_k) \odot s_{k-1} + (1 - \sigma(W_g z_k)) \odot \tanh(W_c z_k), \\
@@ -394,7 +400,7 @@ throughput curve $\SPS(B)$, the scheduler computes:
 \caption{Hardware-Aware Prefix Scheduler (Algorithm~1)}
 \label{alg:scheduler}
 \begin{algorithmic}[1]
-\State \textbf{Input:} confidence scores $\{c_{r,1:\gamma}\}_{r=1}^R$, throughput curve $\SPS(B)$
+\State \textbf{Input:} confidence scores $\{c_{r,1}, \ldots, c_{r,\gamma}\}_{r=1}^R$, throughput curve $\SPS(B)$
 \State \textbf{Output:} verification lengths $\{\ell_r\}_{r=1}^R$
 \For{each request $r$}
   \State $a_{r,j} \gets \prod_{i=1}^j c_{r,i}$ \Comment{Prefix survival probs}
@@ -425,7 +431,7 @@ order of decreasing survival probability, which ensures that extending a
 request's verification length only happens when it adds the globally most
 valuable token. The early stopping condition guarantees the
 non-anticipating property --- the scheduler never extends a request's
-length based on future tokens' survival probabilities that cannot be
+length based on survival probabilities of future tokens, which cannot be
 known at the current step.
 
 A fallback \texttt{StaticScheduler} with fixed-length and static-threshold
@@ -458,7 +464,7 @@ selection and 15-bin ECE computation for calibration quality assessment.
 
 We validate the \UraionSpec implementation through unit tests, gradient
 checks, smoke training, and component-level verification. All experiments
-were conducted on a MacBook Pro 12,1 (Intel i5-5257U, 7.8~GB RAM, CPU only)
+were conducted on a MacBookPro12,1 (Intel i5-5257U, 7.8~GB RAM, CPU only)
 running Ubuntu~24~LTS, targeting the \text{Qwen/Qwen3-0.6B} model.
 
 \subsection{Unit Tests}
@@ -469,24 +475,25 @@ The codebase includes 80 unit tests spanning all major components. Table
 
 \begin{table}[t]
 \centering
+\small
 \caption{Unit test coverage by module. All 80 tests pass.}
 \label{tab:tests}
 \begin{tabular}{lcc}
 \toprule
 Test Suite & Tests & Coverage \\
 \midrule
-Acceptance rule ($\min(1,\; p_t/p_d)$) & 8 & Full: all-accepted, rejected, partial,\\
- & & bonus token, batch independence, \\
+Acceptance rule ($\min(1,\; p_t/p_d)$) & 8 & Full: all, rejected, partial,\\
+ & & bonus token, batch independence,\\
  & & expected length \\
-Markov head & 11 & Vanilla, gated, forward shapes,\\
+Markov \& RNN heads & 15 & Vanilla, gated, RNN forward shapes,\\
  & & gradients, sampling \\
-Scheduler & 12 & Throughput profile, Algorithm~1,\\
- & & static fallback, zero confidence \\
-Shape/gradient checks & 7 & Draft model, losses, confidence head,\\
+Scheduler & 14 & Throughput profile, Algorithm~1,\\
+ & & static fallback, zero conf. \\
+Shape/gradient checks & 10 & Draft model, losses, conf. head,\\
  & & end-to-end cycle \\
-STS calibration & 8 & ECE, temperature fitting, fit/transform,\\
- & & calibrator API \\
-DFlash backbone & 16 & Attention, decoder layer, full backbone,\\
+STS calibration & 8 & ECE, temperature fitting,\\
+ & & fit/transform, calibrator API \\
+DFlash backbone & 17 & Attention, decoder layer, full backbone,\\
  & & GQA, masks, gradient flow, edge cases \\
 Sampling utilities & 8 & Residual, greedy, temperature, gather \\
 \bottomrule
@@ -513,7 +520,7 @@ We verify the lossless speculative decoding acceptance rule
 \paragraph{DFlash backbone.}
 The \DFlash-style backbone passes all shape and gradient checks:
 \begin{itemize}
-  \item Forward pass produces correct output shapes for single/GQA/masked
+  \item Forward pass produces correct output shapes for single/GQA (Grouped Query Attention)/masked
     attention configurations.
   \item Gradient flow verified through attention, decoder layer, and full
     backbone stack.
@@ -524,7 +531,7 @@ The \DFlash-style backbone passes all shape and gradient checks:
 \end{itemize}
 
 \paragraph{Training objective.}
-The three-term loss (Eq.~12) is verified through a smoke training run:
+The three-term loss (Equation~12 of the \DSpark{} paper) is verified through a smoke training run:
 \begin{itemize}
   \item 32 samples from the Capybara dataset, block size $\gamma=4$.
   \item All three loss terms (CE, TV, confidence) are non-negative and
@@ -553,7 +560,7 @@ limitations should be noted:
     CPU-based smoke testing of \text{Qwen3-0.6B}. Full training and
     benchmarking on GPU hardware (e.g., Colab A100) is the natural next
     step.
-  \item \textbf{No vLLM integration}: Unlike the official \DeepSpec\
+  \item \textbf{No vLLM integration}: Unlike the official \DeepSpec{}
     repository, \UraionSpec does not integrate with production serving
     frameworks. We prioritize algorithm clarity over production readiness.
 \end{enumerate}
@@ -564,7 +571,7 @@ limitations should be noted:
 
 All code is publicly available under the MIT license:
 
-\begin{itemize}[leftmargin=*,itemsep=2pt]
+\begin{itemize}
   \item \textbf{GitHub}: \url{https://github.com/arnavprabhu/UraionSpec}
   \item \textbf{HuggingFace}: \url{https://huggingface.co/UraionLabs/UraionSpec}
 \end{itemize}
@@ -604,13 +611,13 @@ colab run --gpu A100 --keep --timeout 28800 \
 \label{sec:conclusion}
 
 We have presented \UraionSpec, a faithful, modular, and verified
-implementation of the \DSpark speculative decoding algorithm. Our
+implementation of the \DSpark{} speculative decoding algorithm. Our
 implementation covers all core algorithmic components---semi-autoregressive
 generation via Markov, gated Markov, and RNN sequential heads, confidence
 prediction with analytical supervision, Sequential Temperature Scaling
-calibration, and the hardware-aware prefix scheduler (Algorithm~1)---plus
-a \DFlash-style backbone with target model KV injection that realizes the
-full parallel drafting architecture described in the original paper.
+calibration, and the hardware-aware prefix scheduler (Algorithm~1).
+We also provide a standalone \DFlash-style parallel backbone module with
+target model KV injection, available for integration into the draft model pipeline.
 
 With 80 passing unit tests, comprehensive documentation, and reproducible
 smoke training and evaluation pipelines, \UraionSpec provides a solid
@@ -622,7 +629,7 @@ inference.
 \paragraph{Future work.}
 Key directions for extending \UraionSpec include full-scale GPU training on
 the Open-PerfectBlend dataset, real engine throughput profiling for the
-scheduler, multi-GPU distributed training via DeepSpeed or FSDP, integration
+scheduler, multi-GPU distributed training via DeepSpeed or FSDP~(Fully Sharded Data Parallelism), integration
 with production serving frameworks (vLLM, llama.cpp), and tree-based
 verification for autoregressive drafters~\cite{spectrindecoding}.
 
@@ -658,7 +665,7 @@ M.~Stern, N.~Shazeer, and J.~Uszkoreit.
 \newblock Blockwise parallel decoding for deep neural machine translation.
 \newblock \textit{arXiv:1811.03115}, 2018.
 
-\bibitem{xia2022medusa}
+\bibitem{xia2024medusa}
 H.~Xia, T.~Ge, S.-M.~Wang, S.-Q.~Chen, F.~Wei, and Z.-Y.~Shao.
 \newblock Medusa: Simple LLM inference acceleration framework with multiple
 decoding heads.
@@ -672,7 +679,7 @@ Y.~Miao, Z.~Bai, Z.~Wang, J.~Zhou, and J.~Jia.
 \bibitem{deepseekspec2026}
 DeepSeek-AI.
 \newblock DeepSpec: Production-grade speculative decoding for LLMs.
-\newblock \texttt{https://github.com/deepseek-ai/DeepSpec}, 2026.
+\newblock \url{https://github.com/deepseek-ai/DeepSpec}, 2026.
 
 \bibitem{spectrindecoding}
 Y.~Zhou, N.~Du, Z.~Zhong, T.~Ji, and Y.~Yang.
@@ -682,17 +689,17 @@ Y.~Zhou, N.~Du, Z.~Zhong, T.~Ji, and Y.~Yang.
 \bibitem{dsparkalphaxiv}
 DeepSeek-AI.
 \newblock DSpark paper at alphaXiv.
-\newblock \texttt{https://www.alphaxiv.org/abs/2026.dspark}, 2026.
+\newblock \url{https://www.alphaxiv.org/abs/2026.dspark}, 2026.
 
 \bibitem{uraionspechf}
 Uraion Labs.
 \newblock UraionSpec on HuggingFace.
-\newblock \texttt{https://huggingface.co/UraionLabs/UraionSpec}, 2026.
+\newblock \url{https://huggingface.co/UraionLabs/UraionSpec}, 2026.
 
 \bibitem{uraionspecgh}
 Arnav Prabhu.
 \newblock UraionSpec on GitHub.
-\newblock \texttt{https://github.com/arnavprabhu/UraionSpec}, 2026.
+\newblock \url{https://github.com/arnavprabhu/UraionSpec}, 2026.
 
 \end{thebibliography}