Role-Aware Neural Convex Divergence Heads for Asymmetric Representation Learning

Technical Summary: Role-Aware Neural Convex Divergence Heads for Asymmetric Representation Learning

1. Problem Statement

Many representation learning tasks involve directed (asymmetric) relations, such as lexical entailment (specific $\to$ general), sentence entailment (premise $\to$ hypothesis), ontology hierarchies (child $\to$ parent), and citation links (citing $\to$ cited). Standard distance functions used in metric learning (Euclidean, cosine, Mahalanobis) are inherently symmetric, failing to capture directionality. Conversely, generic neural scorers (e.g., MLPs on concatenated pairs) can model direction but lack geometric structure, often producing uninterpretable scores that may violate non-negativity or convexity properties. Existing structured approaches like order embeddings or hyperbolic geometry are effective for hierarchies but often function as complete embedding models rather than flexible, plug-in distance heads for arbitrary encoders.

The paper addresses the gap between unstructured asymmetric scorers and rigid geometric models by asking: Can an asymmetric head be both neural and geometrically structured?

2. Methodology

The authors propose the Role-Aware Neural Convex Divergence Head, a plug-in module designed to compute a nonnegative, structured score for ordered pairs $(x, y)$.

Core Components:

• Role Projections: Before computing divergence, the input embeddings $x$ and $y$ are mapped into distinct "source" and "target" roles using learnable linear projections $P_s$ and $P_t$.
  • $u = P_s x$ (Source role)
  • $v = P_t y$ (Target role)
• Neural Bregman Divergence: The head evaluates an Input-Convex Neural Network (ICNN) potential $\phi_\theta$ to compute the Bregman divergence in the projected space:
  $$D_{\theta, P_s, P_t}(x, y) = D_{\phi_\theta}(P_s x, P_t y) = \phi_\theta(u) - \phi_\theta(v) - \nabla\phi_\theta(v)^\top(u - v)$$
• ICNN Constraints: The potential $\phi_\theta$ is parameterized by an ICNN where specific weights are constrained to be non-negative, ensuring convexity. A strongly convex quadratic term ($\lambda \|u\|^2_2$) is added to ensure numerical stability and strict convexity.
• Training Objective: The model is trained using a triplet-style pairwise ranking loss ($L_{rank}$) to ensure positive directed pairs have lower divergence than corrupted negative pairs. A secondary directional margin loss ($L_{dir}$) encourages the forward direction $D(x, y)$ to be smaller than the reverse $D(y, x)$.

3. Key Contributions

The paper makes three primary contributions:

1. Architectural Proposal: A novel plug-in head that combines role-specific projections with neural Bregman divergences. It allows arbitrary encoders to output embeddings that are scored via a structured, nonnegative, asymmetric distance.
2. Theoretical Characterization:
   • Projected-Space Identity: The authors prove that the head maintains non-negativity. Crucially, the zero-divergence condition ($D=0$) occurs when the source projection of the first item equals the target projection of the second ($P_s x = P_t y$), rather than $x=y$. This aligns with the semantic distinctness of source and target roles.
   • Directional Gap Decomposition: In the quadratic special case, the directional gap $G(x, y) = D(x, y) - D(y, x)$ is shown to depend explicitly on the difference between the source-target and target-source projections. This demonstrates that role separation is necessary for directional preference.
   • Local Curvature: The Hessian of the potential provides a mechanism to analyze local convexity, ensuring the learned score behaves as a structured divergence rather than a black-box function.
3. Systematic Empirical Evaluation: The method is evaluated across five benchmarks (HyperLex, WordNet, SICK, SNLI, Gene Ontology) and a large-scale citation stress test (OGBL-Citation2). The study compares the proposed head against symmetric distances, unstructured MLPs, order embeddings, hyperbolic baselines, and plain ICNN-Bregman heads (without role projections).

4. Results

• Directional Accuracy: Across semantic and ontology datasets, the role-aware head consistently outperforms the plain ICNN-Bregman head (without projections) in direction accuracy. For example, on HyperLex, direction accuracy improved from ~38% (plain) to ~90% (role-aware).
• Validity: The proposed head maintains a 0% negative divergence rate, whereas unstructured bilinear scorers frequently produce negative values.
• Ablation Studies: Removing role projections (using shared projections) causes the model to lose directional capability, reverting to performance similar to the plain ICNN. This confirms that role separation is the critical factor for asymmetry, not just the addition of parameters.
• Trade-offs: Increasing the weight of the directional margin improves direction accuracy but can slightly reduce ranking accuracy (R-Acc) on some datasets, indicating a trade-off between distinguishing direction and ranking against negatives.
• Limitations in Large-Scale Graphs: On the OGBL-Citation2 dataset (fixed features, no GNN), specialized symmetric or hyperbolic baselines outperformed the proposed head in ranking metrics (MRR, Hits@K). The authors note that large-scale citation prediction may require stronger input representations or graph encoders.

5. Significance and Claims

The paper positions the Role-Aware Neural Convex Divergence Head not as a universal replacement for all scorers, but as a structured, interpretable plug-in module for tasks where directional relations are central.

• Interpretability: Unlike generic MLPs, this head provides geometric diagnostics (directional gaps, Hessian traces) that allow users to audit why a direction is preferred. The model learns that specific concepts (source) are close to general concepts (target) in a specific projected space.
• Middle Ground: It bridges the gap between rigid geometric models (which assume specific data structures) and flexible neural scorers (which lack geometric guarantees).
• Modest Scope: The authors explicitly state that the method is best suited for scenarios where "geometric validity and interpretation matter." They acknowledge that unstructured MLPs or specialized geometric baselines may still outperform it on pure ranking metrics, particularly when input embeddings already encode the task well or when the evaluation rewards undirected neighborhood quality.

In summary, the work demonstrates that by explicitly modeling source and target roles within a convex divergence framework, one can achieve robust directional discrimination while preserving the mathematical properties (non-negativity, convexity) essential for reliable metric learning.