What Should Feature Distillation Transfer in LLMs? A Task-Tangent Geometry View

TL;DR

This paper introduces a functional geometry perspective for feature distillation in LLMs, emphasizing the transfer of dominant functional directions over raw features, leading to a new method called Flex-KD that improves distillation performance.

Contribution

It proposes a novel functional geometry framework for feature distillation and develops Flex-KD, a parameter-free method that better captures functional contributions, outperforming existing approaches.

Findings

Flex-KD outperforms existing methods across benchmarks.

Effective distillation focuses on dominant functional directions.

Flex-KD handles severe dimension mismatch effectively.

Abstract

Feature-based knowledge distillation aims to transfer intermediate representations from a teacher LLM model to a student. Existing approaches typically rely on direct feature matching or learned projections, implicitly treating representations as objects with intrinsic meaning. However, the relevance of a representation dimension is determined solely by how it affects the model's output. In this work, we propose a functional perspective on feature-based distillation. We characterize knowledge transfer in terms of the teacher's functional geometry, i.e., how its output depends on internal representations, rather than direct representation alignment. This viewpoint reveals that effective distillation need not preserve full high-dimensional features, but instead should retain dominant directions of functional contribution, naturally inducing an effective functional dimension for each task. Building on this framework, we introduce Flex-KD, an architecture-agnostic and parameter-free distillation method that transfers the teacher's functional geometry while matching the student's representational capacity. Extensive experiments across language understanding and generation benchmarks demonstrate that Flex-KD consistently outperforms existing distillation approaches, particularly under severe teacher-student dimension mismatch.