PaperlayerView exact revision on arXiv

Source-linked AI summary

JEDI: JEPA-to-Edge Distillation for Efficient Cropland Segmentation from Satellite Imagery

Kishor Kumar Bhaumik, Nicolas Roque dos Santos, Jia Chen, Evangelos E. Papalexakis

arXiv:2609.07915v1cs.CVcs.LG

TL;DR

Large vision models can provide useful remote-sensing representations but may be too costly for satellite or field-edge deployment, and existing feature-distillation methods often assume similar architectures or stop alignment during task training. JEDI uses two-stage cross-architecture distillation with persistent feature alignment; on CalCROP21, its B0 student reaches 68.0 mIoU with 4.04M parameters, within 2.0 points of the 639M-parameter teacher. The authors state that three limitations bound the findings.

  • Problem

    Large vision models can be too expensive for edge deployment, while feature-distillation methods often assume comparable architectures or discontinue alignment during task training.

    Source: §I · pp. 1–2

  • Method

    JEDI first aligns a compact SegFormer student’s terminal representation to an I-JEPA teacher, then trains segmentation, response distillation, and feature alignment together.

    Source: §I · p. 2

  • Results

    68.0 mIoU: JEDI-B0 uses 4.04M parameters, gains 16.0 points over the standalone student, and is within 2.0 points of the approximately 639M-parameter teacher.

    Source: §I · p. 2, C. Main Results · p. 8

  • Takeaways & Limitations

    The reported CalCROP21 results show that JEDI transfers teacher representations to compact students, with particular benefit under aggressive compression.

    Source: §VI · p. 9, §V · p. 9

    The authors state that three limitations bound their findings.

    Source: §V · p. 9

Abstract

from arXiv · show

Large vision models provide useful representations for remote-sensing segmentation but are often too expensive for deployment at the satellite or field edge. Existing feature-level distillation methods also tend to assume similar teacher and student architectures and often stop feature alignment when task training begins. We introduce JEDI (JEPA-to-Edge Distillation), a two-stage framework that transfers representations from a large I-JEPA Vision Transformer teacher to a compact SegFormer student. First, JEDI aligns the student's terminal representation with the teacher's token space using cross-architecture projection and spatial alignment. It then jointly optimizes supervised segmentation, temperature-scaled response distillation, and persistent feature alignment throughout task adaptation. On CalCROP21, JEDI-B0 achieves 68.0 mean Intersection-over-Union (mIoU) with 4.04M parameters, improving over the standalone student by 16.0 points and coming within 2.0 points of the 70.0 mIoU achieved by the 639M-parameter teacher. We evaluate SegFormer B0, B1, and B2 students with 4.04M, 14.33M, and 28M parameters, respectively. Across all three variants, JEDI consistently outperforms response-, structure-, channel-, and relational-distillation baselines under the same teacher-student setting. These results show that persistent representation alignment is especially valuable under aggressive compression, substantially reducing model size and computation while preserving segmentation performance.

I. INTRODUCTION

JEDI targets efficient cropland segmentation by transferring a large I-JEPA teacher’s representations to compact SegFormer students. On CalCROP21, the 4.04M-parameter B0 reaches 68.0 mIoU, while the framework addresses architectural mismatch and keeps feature alignment active during task adaptation.

I. INTRODUCTION

  • Satellite and field-edge platforms face tighter memory, power, and computation constraints, while retrospective crop maps do not serve applications needing growing-season information.Satellite imagery can support more timely crop mapping, but remote-sensing segmentation also faces spatial variation, intra-class variability, and foreground–background imbalance.Source: §I · p. 1
  • Many feature-distillation methods assume directly matchable teacher and student features, unlike an isotropic ViT token sequence and a hierarchical SegFormer feature pyramid.Their representations differ in embedding dimension, spatial resolution, and topology.Source: §I · p. 2
  • JEDI first aligns the student’s terminal representation with the teacher’s token space, then jointly trains segmentation, response distillation, and persistent feature alignment.The first stage uses no task labels; feature alignment remains active as the student adapts to segmentation.Source: §I · p. 2
  • JEDI selects one NDVI-guided Sentinel-2 observation per tile to reduce inference cost rather than repeatedly evaluating the full temporal stack.The same selection is used during training and inference; the paper frames this as a practical accuracy–compute trade-off, not a replacement for temporal modeling.Source: §I · p. 2
  • 68.0 mIoU: JEDI-B0 uses 4.04M parameters, gains 16.0 points over the standalone student, and comes within 2.0 points of the approximately 639M-parameter teacher.The introduction characterizes the approach as a practical accuracy–compute trade-off across student capacities.Source: Figure 1 · p. 1, §I · p. 2

II. RELATED WORK

Cropland mapping faces challenges from phenology, acquisition conditions, geographic context, and reliable ground truth. Satellite foundation models offer transferable representations, while their capacity motivates efficient transfer to smaller models.

A. Cropland Mapping and Remote-Sensing Segmentation

  • Crop appearance varies with phenology, acquisition conditions, and geographic context, while early mapping, crop rotation, multi-source imagery, and reliable ground truth remain challenges.CalCROP21 combines Sentinel-2 observations over California’s Central Valley with crop labels derived from and refined relative to the USDA Cropland Data Layer.Source: A. Cropland Mapping and Remote-Sensing… · p. 3
  • Remote-sensing segmentation must handle foreground–background imbalance and large intra-class variation beyond the challenges of natural-image segmentation.FarSeg models foreground–scene relations to address these issues.Source: A. Cropland Mapping and Remote-Sensing… · p. 3

B. Satellite Representation Learning and Foundation Models

  • Satellite representation learning adapts self-supervised methods to spectral and temporal observations; SatMAE uses temporal embeddings and spectral positional encodings.The passage reports downstream transfer improvements, including for semantic segmentation.Source: B. Satellite Representation Learning… · p. 3
  • Large geospatial self-supervised models such as Prithvi transfer across Earth-observation tasks, but their capacity increases the need for efficient transfer under compute or memory constraints.I-JEPA predicts target-region representations from context, and JEDI uses it as a training-time teacher for compact SegFormer students rather than deploying the foundation model itself.Source: B. Satellite Representation Learning… · p. 3

C. Efficient Segmentation and Edge Deployment

Efficient segmentation architectures and deployment-oriented compression address resource limits in Earth observation. JEDI instead transfers representations across architectures and retains feature alignment during task adaptation.

C. Efficient Segmentation and Edge Deployment

  • SegFormer pairs a hierarchical Mix Transformer encoder with a lightweight all-MLP decoder and offers model sizes including compact MiT-B0.This reduces deployed-model cost by design rather than relying only on post-training compression.Source: C. Efficient Segmentation and Edge… · p. 3
  • JEDI treats the large teacher as a training-time source of representation knowledge and discards it after distillation, rather than requiring the deployed model to retain its architecture.This contrasts with work that compresses or searches architectures for target hardware,.Source: C. Efficient Segmentation and Edge… · p. 3

D. Knowledge Distillation for Semantic Segmentation

  • Segmentation distillation ranges from softened teacher outputs to structured transfers of pixel relations, channel distributions, intra-class variation, and cross-image relations [11],, [13],,,,.CIRKD transfers global pixel-to-pixel and pixel-to-region relationships across images.Source: D. Knowledge Distillation for Semantic… · p. 4
  • JEDI addresses teacher–student differences in dimensionality, spatial resolution, and topology with learned alignment, and keeps feature alignment active throughout task adaptation.The I-JEPA teacher produces one token sequence, whereas SegFormer produces a multi-resolution feature pyramid.Source: D. Knowledge Distillation for Semantic… · p. 4

III. METHODOLOGY

JEDI trains a compact student on a single selected Sentinel-2 observation, using a temporal stack for acquisition selection. Its two-stage pipeline first aligns student and teacher representations, then adapts the student to segmentation.

A. Problem Formulation and Notation

  • The problem setup uses a multi-temporal Sentinel-2 stack and pixel-wise cropland annotations, but the deployed segmentation model receives one selected three-channel observation.The temporal stack is used only during acquisition selection.Source: A. Problem Formulation and Notation · p. 4

B. NDVI-Guided Composite Selection

  • NDVI-guided selection reduces segmentation-network evaluations from T acquisitions to one; the selected visible bands are cropped and standardized as the network input.The criterion tends to favor low-cloud acquisitions and preserves compatibility with the RGB-pretrained teacher.Source: B. NDVI-Guided Composite Selection · p. 4
  • Stage 1 aligns the trainable student encoder to the frozen teacher without task labels; Stage 2 trains segmentation, response distillation, and persistent feature alignment together.At inference, only the student encoder and decoder are retained.Source: Figure 2 · p. 5

C. Teacher and Student Models

The teacher is a fine-tuned I-JEPA ViT-H/16, while the student uses SegFormer’s hierarchical encoder and lightweight decoder. JEDI aligns the student’s deepest feature with the teacher’s token representation despite differences in width, spatial resolution, and layout.

C. Teacher and Student Models

  • The teacher is a ≈639M-parameter I-JEPA ViT-H/16 fine-tuned for cropland segmentation, and the evaluated SegFormer B0–B2 students span 4.04M to 28M parameters.The teacher encoder and decoder are frozen in subsequent stages.Source: C. Teacher and Student Models · p. 5
  • The student’s terminal stage F(4) is selected because its abstraction level most closely matches the teacher’s single-scale representation.Matching a shallower stage would ask features with less accumulated context to reproduce semantically abstract targets.Source: C. Teacher and Student Models · p. 5

D. Cross-Architecture Alignment Module

  • Directly matching the teacher’s flat token sequence to the student’s channel-first feature map is ill-posed because width, spatial resolution, and layout differ.The alignment module tokenizes the student map, projects tokens to the teacher’s embedding width, and resamples to the teacher’s spatial resolution for position-by-position correspondence.Source: D. Cross-Architecture Alignment Module · p. 5
  • A shallow linear projection keeps the representational burden on the student encoder, while upsampling preserves the teacher’s full spatial grid.The alignment module is discarded after training, so it adds no inference cost.Source: D. Cross-Architecture Alignment Module · p. 5

E. Stage 1: Representation Alignment

JEDI first trains the student encoder and alignment module to reproduce the frozen teacher’s latent geometry without task supervision. It then adapts the aligned student with segmentation, response distillation, and feature alignment jointly, keeping representational supervision active throughout adaptation.

E. Stage 1: Representation Alignment

  • Stage 1 trains the student encoder and projection jointly against frozen teacher targets without task supervision, excluding the decode head from updates.Squared error aligns both representation direction and magnitude; the encoder and projection are checkpointed to initialize Stage 2.Source: E. Stage 1: Representation Alignment · pp. 5–6

F. Stage 2: Tri-Objective Joint Refinement

  • Stage 2 warm-starts the student encoder and alignment module from Stage 1, retains the decode head’s ImageNet initialization, and trains all three components.The objective combines segmentation, response distillation, and persistent feature alignment.Source: F. Stage 2: Tri-Objective Joint… · p. 6
  • Temperature-scaled response distillation transfers the teacher’s inter-class uncertainty, which may be informative at field boundaries and spectrally ambiguous regions.Teacher logits are resampled to the student lattice before matching softened probabilities.Source: F. Stage 2: Tri-Objective Joint… · p. 6
  • During Stage 2, JEDI retains feature alignment as an explicit constraint so task adaptation continues to preserve similarity to the teacher representation.The authors identify removing this constraint (γ = 0) as an ablation to isolate its effect.Source: F. Stage 2: Tri-Objective Joint… · p. 6
  • At inference, only the student encoder and decoder remain; the teacher and alignment module are discarded.The training procedure reuses encoder outputs to construct its loss terms.Source: F. Stage 2: Tri-Objective Joint… · p. 6

A. Experimental Setup

The experiments evaluate binary cropland segmentation on CalCROP21 using single NDVI-selected Sentinel-2 composites. They report segmentation quality and model efficiency, with results averaged over five seeds.

A. Experimental Setup

  • CalCROP21 evaluation uses binary cultivated-versus-non-cultivated segmentation on 448 × 448 single-composite tiles from multi-temporal Sentinel-2 data.The tiles are selected using NDVI and the study adopts the train/test partition of [29].Source: A. Experimental Setup · p. 6
  • Training uses AdamW with cosine-annealed learning rate, 30 Stage 1 epochs, 20 Stage 2 epochs, batch size 8, temperature 2, and loss weights (1.0, 0.5, 0.5).The models are trained on a single NVIDIA RTX A6000.Source: A. Experimental Setup · pp. 6–7
  • Segmentation quality is measured by mIoU and macro F1, while efficiency is measured by parameter count, GMACs, and single-A6000 inference latency.Results are reported as mean ± standard deviation over five independent seeds.Source: A. Experimental Setup · p. 7

B. Baselines

The section compares JEDI with established distillation objectives and reports results across three student capacities. JEDI improves on standalone students and prior distillation baselines while bringing the smallest student close to the teacher.

B. Baselines

  • The comparison adapts KD, SKD, CWD, and CIRKD to the same frozen teacher–student pair, routing feature-level objectives through shared alignment.Source: B. Baselines · p. 7

C. Main Results

  • 68.0±0.4 mIoU and 80.0±0.3 F1: JEDI-B0 beats standalone MiT-B0 by 16.0 mIoU and 14.0 F1 points, and CIRKD by 2.1 and 2.3 points.Source: C. Main Results · p. 7, Table II · p. 7
  • MiT-B1 and MiT-B2 reach 69.0 and 69.5 mIoU, with JEDI-over-standalone gains declining from 15.0 to 13.5 points as capacity rises.Source: C. Main Results · p. 7
  • Against the 639.12M-parameter teacher at 70.0 ± 0.3 mIoU, JEDI-B0 reaches 68.0±0.4 mIoU with 4.04M parameters, approximately 158× fewer.Source: C. Main Results · pp. 7–8

D. Ablation Studies

Ablations indicate that representation-alignment initialization and continued feature alignment each contribute to JEDI, while response distillation adds complementary information. In the evaluated teacher-choice setting, I-JEPA produces the strongest distilled student.

D. Ablation Studies

E. Choice of Teacher

F. Efficiency Analysis

JEDI-B0 substantially reduces computational cost while retaining segmentation performance near the teacher. Qualitative results and ablations support persistent alignment, while the findings remain bounded by the study’s dataset, task, and acquisition choices.

F. Efficiency Analysis

  • 68.0 mIoU is retained by JEDI-B0 while reducing GMACs from 557.5 to 5.8 and parameter count by more than two orders of magnitude versus the teacher.GMACs are approximate lower bounds; attention, GELU, and softmax operations are not counted.Source: F. Efficiency Analysis · p. 8, Table V · p. 9

G. Qualitative Results

  • On three randomly selected test tiles, JEDI-B0 more closely tracks the teacher and recovers coherent field geometry than prior distillation methods.The prior methods recover coarse field layout but leave larger misclassified regions, especially at field boundaries and in spectrally ambiguous interiors.Source: G. Qualitative Results · p. 9

V. DISCUSSION AND LIMITATIONS

  • MiT-B0 gains 16.0 points over standalone training, while B1 and B2 gain less; omitting Stage 2 alignment erodes Stage 1’s representation initialization.Response distillation adds a smaller gain, and the ablations support combining it with representation alignment under aggressive compression.Source: §V · p. 9
  • The findings cover binary cropland segmentation on CalCROP21 in California’s Central Valley; broader geographic or sensor generalization remains untested.The NDVI-guided selection also reduces the temporal stack to one acquisition, discarding trajectory information that could distinguish spectrally similar crops.Source: §V · p. 9

VI. CONCLUSION

  • Across student capacities, JEDI reduces parameter count and GMACs while retaining most of the teacher’s segmentation performance.The conclusion reports that both pre-adaptation representation alignment and its persistence during adaptation contribute to final performance, with removing persistent alignment causing the largest degradation.Source: §VI · p. 9

Related papers

Loading 2609.07915v1…