Source-linked AI summary
Next-Embedding Prediction Makes Strong Vision Learners
Sihan Xu, Ziqiao Ma, Wenhao Chai, Xuweiyi Chen, Weiyang Jin, Joyce Chai, Saining Xie, Stella X. Yu
TL;DR
The paper asks whether causal generative pretraining can yield strong visual learners without conventional representation-learning machinery. NEPA predicts future patch embeddings from past ones with a causal Transformer, and shows effective transfer across ImageNet-1K classification and ADE20K segmentation. The method offers a simple predictive alternative that may extend across modalities, within the demonstrated scope.
Problem
Visual self-supervised learning commonly trains representations for downstream modules, motivating whether minimal causal generative prediction can instead produce strong visual learners.
Method
NEPA trains a causal Transformer to predict future patch embeddings from previous ones using stop-gradient targets and no pixel decoder, discrete tokenizer, contrastive loss, or auxiliary encoder.
Results
NEPA transfers effectively across ImageNet-1K classification and ADE20K semantic segmentation, reaching 83.8% and 85.3% top-1 accuracy with Base and Large models.
Takeaways & Limitations
Next-embedding prediction can learn transferable visual representations with a simpler single-stream autoregressive training paradigm using only ImageNet-1K self-supervised pretraining.
Takeaways & Limitations
The model struggles with reflections, shading, shadows, and scenes containing many small or overlapping objects, possibly reflecting limitations of ImageNet training data.
Abstract
from arXiv · showhide
Inspired by the success of generative pretraining in natural language, we ask whether the same principles can yield strong self-supervised visual learners. Instead of training models to output features for downstream use, we train them to generate embeddings to perform predictive tasks directly. This work explores such a shift from learning representations to learning models. Specifically, models learn to predict future patch embeddings conditioned on past ones, using causal masking and stop gradient, which we refer to as Next-Embedding Predictive Autoregression (NEPA). We demonstrate that a simple Transformer pretrained on ImageNet-1k with next embedding prediction as its sole learning objective is effective - no pixel reconstruction, discrete tokens, contrastive loss, or task-specific heads. This formulation retains architectural simplicity and scalability, without requiring additional design complexity. NEPA achieves strong results across tasks, attaining 83.8% and 85.3% top-1 accuracy on ImageNet-1K with ViT-B and ViT-L backbones after fine-tuning, and transferring effectively to semantic segmentation on ADE20K. We believe generative pretraining from embeddings provides a simple, scalable, and potentially modality-agnostic alternative to visual self-supervised learning.
1. Introduction
NEPA asks whether causal generative prediction in embedding space can produce strong visual learners, replacing conventional representation-learning objectives with next-patch-embedding prediction. A simple Transformer pretrained only with this objective transfers effectively to classification and semantic segmentation.
- 1. Introduction: The approach follows a broader shift from visual feature extraction toward predictive models that learn conditional dependencies, paralleling generative pretraining in language.Earlier visual generative approaches faced long-sequence or semantic-alignment limitations, motivating prediction in embedding space.
- 1. Introduction: NEPA predicts each future patch embedding from preceding embeddings, using causal prediction and stop-gradient rather than pixel reconstruction or discrete visual tokens.The objective also avoids contrastive augmentations, negative pairs, momentum encoders, and task-specific architectural components.
- 1. Introduction: NEPA also transfers effectively to ADE20K semantic segmentation, showing that its learned visual features support dense prediction beyond image classification.The introduction characterizes these results as strong transfer from a purely predictive model.
2. Related Work
Related visual self-supervised methods rely mainly on invariance learning, reconstruction, or representation prediction, each with characteristic architectural or data requirements. NEPA instead applies a simple causal next-token-style objective directly in continuous visual embedding space.
- 2. Related Work: Contrastive and self-distillation methods learn invariances across augmented views, while reconstruction methods recover masked pixels or tokens using lightweight decoders.Contrastive approaches often require large batches or memory banks, whereas reconstruction approaches introduce decoding objectives.
- 2. Related Work: Generative visual pretraining has used discrete pixel tokens, continuous pixel targets, or compressed tokenizer codes to make image sequences suitable for Transformer prediction.iGPT learned transferable features but suffered from long sequences and weak semantic alignment.
- 2. Related Work: Representation-prediction methods forecast latent causes or embeddings, but JEPA-style systems typically use auxiliary encoders and prediction heads and predict targets in parallel rather than causally.The passage contrasts these designs with the causal next-token objective used by generative language models.
3. Method
NEPA divides images into patch embeddings and trains a causal Transformer to predict the next embedding from prior context. It uses normalized negative-cosine similarity with stop-gradient targets, while downstream heads adapt the pretrained model to classification and segmentation.
- 3. Method: The model converts an image into a sequence of T patch embeddings and autoregressively predicts each next embedding conditioned on the preceding sequence.This is implemented with a shared embedding encoder and autoregressive predictor using causal attention.
- 3. Method: Predicted and target embeddings are normalized and compared with negative cosine similarity, encouraging semantic alignment without explicit input reconstruction.The target embedding comes directly from the encoder for the next patch.
- 3. Method: Stop-gradient is applied to target embeddings to avoid degenerate solutions during optimization.Without stop-gradient, the pretraining loss collapses to −1 and all embeddings become identical.
- 3.2. Model Architecture: The architecture uses a Conv2d patch embedder, pre-norm Transformer, LayerNorm, causal masking, and no separate decoder for prediction.RoPE, LayerScale, SwiGLU, and QK-Norm are included as stability and scalability components orthogonal to the core framework.
- 3.3. Downstream Tasks: Classification uses the final embedding with a linear head, while segmentation uses UPerNet multi-scale features and bidirectional attention during fine-tuning.These heads provide task-specific supervision only at the downstream stage.
4. Experiments
Experiments show that NEPA depends on causal prediction design choices for stable training and transfers effectively to classification and segmentation, while scaling improves accuracy without overfitting.
- Core algorithm design: Removing causal masking reduces accuracy from 76.8% to 73.6% top-1 after 50k pretraining steps, while removing the autoregressive shift causes fine-tuning divergence.The shift prevents an identity-mapping objective, and causal masking preserves the autoregressive prediction setup.
- Core algorithm design: Removing stop-gradient causes training collapse, with the loss reaching −1 and all embeddings becoming identical.Stop-gradient is applied to ground-truth embeddings to prevent collapse.
- Core algorithm design: 0% input masking achieves 78.2% top-1 accuracy, versus 76.4% with 40% masking and 75.7% with 60% masking after 100k pretraining steps.The reported trend contrasts with masked image modeling, because causal embedding prediction already supplies a nontrivial learning signal.
- Orthogonal components: LayerScale, RoPE, QK-Norm, and SwiGLU are evaluated as architectural additions, with their combination reported as the best configuration and RoPE providing a significant accuracy improvement.LayerScale stabilizes pretraining but can slow supervised convergence, while QK-Norm improves stability and final performance.
- Scaling behavior: Accuracy steadily improves with increased training steps, with no observed overfitting under extended compute budgets.The scaled models support subsequent attention, embedding, classification, and segmentation analyses.
- Classification results: 83.8% and 85.3% top-1 accuracy are achieved on ImageNet-1K by the Base and Large models, respectively, with competitive performance from a single-stream embedding-prediction formulation.The method uses no task-specific head, decoder, reconstruction target, auxiliary loss, or multi-branch architecture during pretraining.
- Semantic Segmentation: 48.3% and 54.0% mIoU are reached on ADE20K by the Base and Large models, respectively, after ImageNet-1K pretraining.The models transfer to dense prediction despite using neither a decoder nor a pixel-level pretraining objective.
5. Quantitative Results
NEPA’s analyses examine whether next-embedding prediction uses global, semantically organized visual context. Attention and similarity maps indicate object-centric dependencies rather than purely local or texture-based matching.
- Classification Results: Table 2 compares self-supervised frameworks on ImageNet-1K classification across Base and Large model scales.The caption specifies grouped model-scale blocks and effective pretraining epochs.
- Semantic Segmentation: Table 3 compares different pretraining methods on ADE20K semantic segmentation using mIoU.The supplied caption identifies the comparison scope but does not provide cell values.
- Attention Map Analysis: NEPA’s attention is often long-ranged and object-centric, concentrating on semantically related object regions rather than uniformly or locally attending.The analyzed examples show attention to other object parts and nearby informative structures across categories and viewpoints.
- Embedding Similarity Analysis: Predicted embeddings are most similar to patches from the same object or semantic region, while unrelated backgrounds show lower similarity.This organization emerges without explicit labels or region annotations.
- Embedding Similarity Analysis: Distinctive query patches produce similarities to distant, occluded, or nonlocal object parts, indicating extrapolation of object-level structure beyond local texture.Homogeneous backgrounds yield more diffuse but still coherent similarity patterns, reflecting uncertainty.
6. Conclusion and Future Work
The conclusion presents next-embedding prediction as a simple embedding-space analogue of causal next-token prediction. It reports transferable visual representations, modality-agnostic potential, and possible extensions toward image generation, alongside deployment cautions.
- Conclusion: NEPA learns transferable visual representations with ImageNet-1K self-supervised pretraining while avoiding handcrafted pretext-task design choices.The conclusion describes a single forward pass and no decoder, with classification and segmentation transfer matching or slightly outperforming mainstream methods.
- Modality-agnostic potentials: Embedding-space autoregression offers a common objective across modalities, with embeddings serving as a shared representational currency.The paper connects NEPA to tied input-output embeddings increasingly used in language models.
- Generative potentials: A suitable image decoder or diffusion-based generator could extend NEPA toward image synthesis or editing, but this remains future work.The proposed direction would bridge representation learning and generation within one model.
- Broader Impacts: NEPA may inherit dataset biases or be misused to create misleading visual content, requiring care in data selection and deployment safeguards.The paper specifically calls for monitoring alongside appropriate safeguards.
B. Methodology Comparisons
The comparisons distinguish NEPA from contrastive, pixel-target, and JEPA-style approaches by placing causal prediction directly in continuous embedding space. NEPA retains latent prediction while simplifying the architecture.
- Relation to CPC: Unlike CPC, NEPA directly trains the transformer predictor to regress the next embedding without negatives or a contrastive head.CPC primarily trains an encoder, while its autoregressive module mainly aggregates context.
- Relation to GPT, iGPT and AIM: NEPA predicts continuous embeddings rather than predefined pixel-space targets such as discrete pixel tokens or continuous patch values.The method uses representation-space supervision through a learned, moving embedding target.
- Relation to JEPA: Compared with JEPA, NEPA keeps latent prediction but uses one embedding layer and an autoregressive transformer without asymmetric branches or an extra head.The supplied Figure 6 caption identifies the JEPA–NEPA comparison, while the prose states the architectural simplification.
C. Implementation Details
The implementation uses standard computer-vision ecosystems and recipes for large-scale pretraining and transfer. ImageNet-1K is accessed through Hugging Face datasets, while ADE20K follows MMSegmentation pipelines.
- Software Stack: NEPA was prototyped with timm and later implemented at scale using Hugging Face Transformers, Datasets, and Evaluate.The implementation is based mainly on the Transformers ViT code and adapted official training examples.
- Data and Pipelines: Training uses timm utilities for Mixup and image transformation, with ImageNet-1K accessed through Hugging Face Datasets.ADE20K experiments use standard MMSegmentation recipes and data-processing pipelines.
D. Training Details
NEPA pretraining uses standard learning-rate scaling, while training dynamics are examined through loss curves and evolving attention maps.
- The pretraining procedure uses linear learning-rate scaling, setting lr = base lr × B/256 for global batch size B.
- Table 4 provides the pretraining settings used by the models.
- Pretraining dynamics are assessed using loss curves and intermediate NEPA-L attention maps to illustrate training stability and attention evolution.
E.2. Limitations and Failure Cases
NEPA shows a mismatch between shallow linear-probe performance and broader predictive capacity, while failure cases cluster around difficult physical interpretation and complex scenes.
- NEPA performs poorly under standard linear probing because the evaluated output remains close to shallow post-embedding features rather than exposing the predictor’s full capacity.
- The reported limitations concern both shallow linear probing and complex visual reasoning, rather than a failure of the pretrained predictor across all evaluations.
- NEPA struggles with reflections, shading, shadows, and scenes containing many small or overlapping objects, producing uncertain or inconsistent predictions.
- The authors hypothesize that ImageNet’s limited diversity contributes to these reasoning failures and may be addressed by scaling to more diverse datasets.
E.5. Reproduced Results without SwiGLU
Removing SwiGLU preserves NEPA’s performance, while attention and embedding visualizations show semantically aligned predictions across diverse ImageNet classes.
- The no-SwiGLU comparison uses tuned learning-rate and layer-wise decay settings to report the best reproduced configuration.
- Similar results are reproduced for NEPA-B without SwiGLU, indicating that the main gains come from the predictive architecture rather than this activation choice.
- Across eight ImageNet classes, NEPA attends to semantically related regions and assigns high predicted–true embedding similarity to patches on the same object.
- The reported visualization set includes query maps, attention maps over seen patches, and cosine-similarity maps between predicted and true embeddings.