Recursive KL Divergence Optimization: A Dynamic Framework for Representation Learning

Abstract: We propose a generalization of modern representation learning objectives by reframing them as recursive divergence alignment processes over localized conditional distributions While recent frameworks like Information Contrastive Learning I-Con unify multiple learning paradigms through KL divergence between fixed neighborhood conditionals we argue this view underplays a crucial recursive structure inherent in the learning process. We introduce Recursive KL Divergence Optimization RKDO a dynamic formalism where representation learning is framed as the evolution of KL divergences across data neighborhoods. This formulation captures contrastive clustering and dimensionality reduction methods as static slices while offering a new path to model stability and local adaptation. Our experiments demonstrate that RKDO offers dual efficiency advantages approximately 30 percent lower loss values compared to static approaches across three different datasets and 60 to 80 percent reduction in computational resources needed to achieve comparable results. This suggests that RKDOs recursive updating mechanism provides a fundamentally more efficient optimization landscape for representation learning with significant implications for resource constrained applications.

• The paper's main contribution is a dynamic framework that recursively refines supervisory distributions via EMA, achieving about 30% lower loss than static methods.
• It employs a practical implementation with student-teacher models and memory banks to update evolving conditional distributions during training.
• Experiments on CIFAR-10, CIFAR-100, and STL-10 show RKDO attains competitive performance with up to 80% fewer training epochs, optimizing resource efficiency.

Recursive KL Divergence Optimization (RKDO) (2504.21707) is a framework for representation learning that reframes objectives like contrastive learning, clustering, and dimensionality reduction as processes of recursively aligning localized conditional distributions. It builds upon the Information Contrastive Learning (I-Con) framework [alshammari2025icon], which views many representation learning methods as minimizing the KL divergence between a supervisory distribution $p(j|i)$ and a learned distribution $q(j|i)$ over neighborhood pairs $(i, j)$. RKDO argues that this process is not static but inherently dynamic and recursive.

The core idea behind RKDO is that the supervisory distribution $p^{(t)}(\cdot|i)$ for a data point $x_i$ at iteration $t$ is not fixed but recursively depends on the learned distribution $q^{(t-1)}(\cdot|i)$ from the previous iteration $t-1$. This creates a dynamic system where the targets themselves evolve based on what the model has learned.

Core Mechanism

RKDO formalizes this dynamic relationship with the following recursive update for the supervisory distribution $p^{(t)}$ based on the learned distribution $q^{(t-1)}$:

$$p^{(t)}(\cdot|i) = (1 - \alpha) \cdot p^{(t-1)}(\cdot|i) + \alpha \cdot q^{(t-1)}(\cdot|i)$$

Here, $p^{(t)}(\cdot|i)$ and $q^{(t)}(\cdot|i)$ are distributions over neighbors $j$ conditioned on point $i$. The parameter $\alpha \in (0, 1]$ controls the influence of the previous learned distribution $q^{(t-1)}$ on the new target distribution $p^{(t)}$. The learned distribution $q^{(t)}(\cdot|i)$ is computed using the current model parameters $\phi^{(t)}$ and a temperature parameter $\tau^{(t)}$, typically based on dot products of embeddings:

$$q^{(t)}(j|i) = \frac{\exp(f_\phi^{(t)}(x_i) \cdot f_\phi^{(t)}(x_j)/\tau^{(t)})}{\sum_{k\neq i}\exp(f_\phi^{(t)}(x_i) \cdot f_\phi^{(t)}(x_k)/\tau^{(t)})}$$

The temperature $\tau^{(t)}$ can also be dynamic, decaying over time: $$\tau^{(t)} = \tau^{(0)} \cdot (1 - \beta \cdot \frac{t}{T})$$, where $\beta$ controls the decay rate over total iterations $T$.

The loss function optimized at each iteration $t$ is the average KL divergence over all data points $i$:

$$L^{(t)} = \frac{1}{n}\sum_{i=1}^{n} D_{KL}(p^{(t)}(\cdot|i) \| q^{(t)}(\cdot|i))$$

Comparison with Previous Work

While the recursive update structure using exponential moving averages (EMA) has been utilized in methods like Temporal Ensembling [laine2017temporal], Mean Teacher [tarvainen2017mean], MoCo [he2020momentum], BYOL [grill2020bootstrap], and DINO [caron2021emerging], these methods typically apply EMA to model weights or per-sample predictions/features. RKDO's novel contribution lies in applying this recursion to the entire field of conditional distributions $p(\cdot|i)$ across all data points $i$. This means the target distribution for each sample is continuously updated based on the model's previous output distribution for that sample.

Practical Implementation

Implementing RKDO requires managing the evolving target distributions $p^{(t)}(\cdot|i)$ for all data points $i$. For large datasets, storing the full $N \times N$ or $N \times K$ (where $K$ is neighborhood size) distribution matrix $p^{(t)}$ can be memory-intensive. A practical implementation strategy, analogous to MoCo/BYOL/DINO, likely involves:

1. Maintaining a student model ($\phi$) and a teacher model ($\phi_{\text{teacher}}$), where $\phi_{\text{teacher}}$'s weights are an EMA of $\phi$'s weights.
2. Using a memory bank storing features (e.g., from the teacher model) of previous batches. This bank defines the set of potential neighbors $j$ for computing $q(j|i)$ and $p(j|i)$.
3. Storing the target distributions $p(\cdot|i)$ for all $N$ training samples over the indices of the memory bank. This requires an $N \times \text{MemoryBankSize}$ tensor for $p_{\text{target\_field}}$.
4. In each training iteration, for a batch of samples:
   • Compute student embeddings $\phi(x_i)$ and teacher embeddings $\phi_{\text{teacher}}(x_i)$.
   • Compute $q_{\text{student}}(j|i)$ and $q_{\text{teacher}}(j|i)$ for $i$ in the batch and $j$ in the memory bank using the current temperature $\tau^{(t)}$.
   • Retrieve the previous target distribution $p^{(t-1)}(\cdot|i)$ for samples in the batch from $p_{\text{target\_field}}$.
   • Compute the new target $p^{(t)}(\cdot|i)$ for the batch: $$p^{(t)}(\cdot|i) = (1-\alpha) p^{(t-1)}(\cdot|i) + \alpha q_{\text{teacher}}^{(t-1)}(\cdot|i)$$. Note: The paper's equation uses $q^{(t-1)}$, which would mean $q$ from the previous iteration. Using $q_{\text{teacher}}$ from the current iteration's teacher seems a practical approximation aligned with related works.
   • Update $p_{\text{target\_field}}$ with the newly computed $p^{(t)}(\cdot|i)$ for the batch indices. Detaching this update prevents gradients from flowing back through $\alpha$ into previous $p$ values.
   • Compute the loss $$D_{KL}(p^{(t)}(\cdot|i) \| q_{\text{student}}^{(t)}(\cdot|i))$$ averaged over the batch.
   • Perform gradient descent on the student model $\phi$.
   • Update the teacher model $\phi_{\text{teacher}}$ via EMA of $\phi$.
   • Update the memory bank.

The experiments in the paper used a ResNet-18 backbone with a projection head, trained on CIFAR-10, CIFAR-100, and STL-10 datasets with standard augmentations. They used a recursion depth of 3 (how many previous $q$ values influence $p$, effectively related to the EMA $\alpha$), $\tau=0.5$, $\beta=0.1$. The code is stated to be available on GitHub.

Key Findings and Practical Implications

The experiments highlight RKDO's "dual efficiency advantages":

1. Optimization Efficiency: RKDO consistently achieved approximately 30% lower loss values compared to the static I-Con approach across all datasets and training durations (Table 1). This suggests RKDO's recursive target update smooths the optimization landscape, allowing the model to reach lower minima.
   • Practical Implication: Lower loss values often correlate with better-learned representations. Achieving this with the same or less effort is a significant gain.
2. Computational Resource Efficiency: RKDO required 60-80% fewer computational resources (training epochs) to achieve performance comparable to longer I-Con training (Tables 2 & 3, Figures 3 & 4). The authors quantify a "resource unit" as an optimizer update step, noting the FLOPs per step are virtually identical for RKDO (depth 3) and I-Con (depth 1). Thus, reduced steps translate directly to reduced wall-clock time, energy, and FLOPs.
   • Concrete Examples: On CIFAR-100, RKDO at 2 epochs achieved performance comparable to or superior to I-Con at 5 epochs, using 60% fewer resources. On STL-10, RKDO at 2 epochs matched I-Con at 5 epochs. On CIFAR-10, RKDO at 1 epoch reached 76% of I-Con's 5-epoch performance with 80% fewer resources.
   • Practical Implication: This is transformative for resource-constrained environments (e.g., edge devices, mobile AI) and applications requiring rapid retraining or learning from limited data. It means potentially deploying models faster and at lower computational cost.

Trade-offs and Dynamics

The paper notes that RKDO shows strong advantages in early training epochs but might see its lead diminish or reverse with extended training on some datasets (e.g., CIFAR-10 after 5-10 epochs). This is likened to "unbounded recursion" in programming, where continuous refinement without a "base case" might lead to overspecialization on the training data. The recursive nature of $p^{(t)}$ becoming more aligned with $q^{(t-1)}$ means the target distribution becomes increasingly specific to the model's own output from the previous step.

Theoretical Analysis

The paper provides a theoretical analysis showing that, under ideal assumptions (infinite model capacity, exact optimization), RKDO achieves linear-rate convergence $L^{(t)} \leq (1-\alpha)^t L^{(0)}$. This formalizes why RKDO can achieve lower loss values. In practice, finite capacity and imperfect optimization mean the loss converges to a non-zero $L_\star$, but still at a linear rate $O((1-\alpha)^t)$ towards that optimum. This theory provides guidance on the role of $\alpha$ (larger $\alpha$ can speed convergence but might affect the minimum $L_\star$) and justifies early stopping when the loss curve flattens.

Implementation Considerations and Deployment

• Parameters: The parameters $\alpha$ (controlling the influence of previous $q$ on $p$), $\beta$ (temperature decay), and potentially the recursion depth (though the paper mainly focuses on depth 3) are crucial tuning knobs. Their settings will likely influence the trade-off between early optimization speed and long-term generalization. Adaptive schedules for $\alpha$ could be beneficial, perhaps starting higher for rapid initial learning and decreasing later.
• Memory: As discussed, storing the $N \times \text{BankSize}$ tensor for $p_{\text{target\_field}}$ is the main memory overhead compared to standard contrastive methods. For very large datasets, this might become a limiting factor depending on available GPU/CPU memory.
• Compute: The per-step computational cost is shown to be comparable to static methods once the memory bank/teacher mechanism is in place. The gain comes purely from requiring fewer steps.
• Deployment: RKDO primarily affects the training phase. Once trained, the model $f_\phi$ is deployed like any other learned representation model, with no extra computational cost during inference. Its benefit is enabling faster, cheaper training to potentially achieve a good model.

Limitations

The paper acknowledges limitations including sensitivity to hyperparameters (like $\alpha$), the focus on relatively short training durations (up to 10 epochs), and evaluation primarily on smaller image datasets. Further research is needed to understand RKDO's behavior on larger scales and longer training regimes, and to develop robust parameter tuning strategies.

Conclusion

RKDO offers a new perspective on representation learning objectives, viewing them as recursive processes. By applying EMA-style recursion to the target conditional distributions for each sample, it demonstrates significant empirical advantages in optimization efficiency (30% lower loss) and computational resource reduction (60-80% fewer epochs) compared to static approaches like I-Con on tested datasets. This makes RKDO a promising framework for applications where training speed and cost are critical, although managing its dynamic optimization landscape and potential for overspecialization with extended training may require further attention.